PuzzleFusion++: Auto-agglomerative 3D Fracture Assembly by Denoise and Verify

We thank the reviewer for the construction comments and the questions. We will address the questions and concerns raised by the reviewer below.

Q1: Speed is significantly slower than some methods

We thank you for pointing out the speed limitation. We will discuss this question from two perspectives.

• The use case of shape assembly probably does not require a fast solver. As discussed in the literature on 3D shape assembly, the main use cases are archeological artifact reconstruction, forensic object reassembly, and protein structure analysis for drug discovery. In these applications, the users care much more about the final assembly quality than the running speed. Jigsaw, the strongest baseline, also needs a rather long running time to solve for one object.

• Employ more advanced diffusion-based generative models to accelerate the inference. The relatively slow speed is mainly due to the denoiser requiring multiple denoising steps. Recent advances in diffusion models, flow matching, consistency models, etc. (such as [A], [B], and [C]) indicate that there is still significant room for improving the sampling speed with fewer sampling steps. Our current setup uses 20 denoising steps, suggesting the potential for up to a 20× speed improvement. This is a promising direction for future work, and we will include the discussion in the paper.

Q2: Can the proposed method solve the reassembly task when the number of fractures is big (eg, 100)

We followed the same setting as the baselines and focused on ≤20 fractures. While we achieved SOTA performance for 20 fractures, the results are far from perfect. With 100 fractures, local geometric ambiguities would become significantly more severe, amplifying the issues observed in our failure cases section (Sec.4.4) and likely resulting in very poor performance. Therefore, we did not investigate cases with more fracture pieces in our experiments. Handling more fractures is an important direction for future work.

Q3: How well can one expect the proposed method to work on unseen objects?

We thank you for the question. We tested our method on the Artifact subset (unseen) of the Breaking Bad dataset, as shown in the main table under "Trained on the everyday subset, tested on the artifact subset." The results are reasonable compared to our main baseline, Jigsaw.

However, PuzzleFusion++ experiences a performance decline across all metrics compared to Jigsaw on the unseen objects. This is because PuzzleFusion++ learns global spatial priors by the diffusion model that simultaneously solves the arrangements of all the pieces. The global priors are effective for everyday objects but struggle to generalize to the different categories in the Artifact subset. In contrast, Jigsaw focuses on local geometry learning, making it less sensitive to major category changes.

To further investigate, we finetuned our denoiser on the Artifact subset and provide qualitative results in Figure 17 and quantitative results in Table 9. With less than 20% of the training iterations, the model can be adapted to previously unseen object categories with good performance. Together with the generalization results in Table 1, we believe our method has reasonable ability to handle unseen objects

Software/Library for rendered images:

We render our results using BlenderToolBox[D]. We will release our rendering code together with other parts of the project code.

[A] Lu, C., & Song, Y. (2024). Simplifying, Stabilizing and Scaling Continuous-Time Consistency Models. arXiv preprint arXiv:2410.11081.

[B] Lipman, Y., Chen, R. T., Ben-Hamu, H., Nickel, M., & Le, M. (2022). Flow matching for generative modeling. arXiv preprint arXiv:2210.02747.

[C] Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, 26565-26577.

[D] Derek Liu. BlenderToolBox. Available at: https://github.com/HTDerekLiu/BlenderToolbox

Thank you for your feedback and support of our work! We are glad our responses addressed your concerns.

We thank you for the construction comments and the questions.

W1 & Q1.1: Could PuzzleFusion be adapted to handle 3D tasks as a baseline for comparison?

PuzzleFusoin is specifically designed for 2D jigsaw puzzles with simple polygonal shapes, where a 1D chain of 2D coordinates represents each puzzle piece. The approach (including its puzzle piece encoder and the denoising network) cannot be directly used for the 3D shape assembly task, where each piece is a set of unordered 3D points.

W1 & Q1.2: What distinguishes PuzzleFusion++ from PuzzleFusion, particularly in features that enhance its suitability for 3D tasks?

The 3D fracture assembly task poses several challenges: i) The 6-DoF solution space is more complicated than the 3-DoF solution space of 2D spatial puzzles; ii) Unlike 2D puzzle solvers that mainly leverage image semantics or polygon structures, 3D fracture assembly requires a deep understanding of fracture surfaces; iii) There can be many small 3D fragments, further complicating the problem.

PuzzleFusion++ proposes new designs to address the above challenges. Specifically, our VQ-VAE with PointNet++ can encode fine details of fracture surfaces into a set of local latents, boosting the shape understanding of the denoising network; Our auto-agglomerative framework verifies assembly results and composes confident small pieces into larger ones to facilitate future iterations, which effectively improves the assembly success rate; We also design tailored components for the diffusion model, including the noise scheduler and the denoising transformer with local shape encodings as condition.

Q2.1: The matching from Jigsaw is not perfect, is there a scheme to handle such situation?

As shown in the Table 6, the verifier's performance is not optimal. We cannot address most of the errors made in Jigsaw point matching. However, our verifier leverages a Transformer to incorporate all pairwise matching information between fragments. This design enables the model to reason globally across multiple pairs, potentially overcoming local matching errors produced during Jigsaw matching.

Q2.2: Upper bound rotation error is still very high (34 degree error)

Thank you for this great observation. We agree on this point — the high rotation error with even the GT matching indeed indicates the limitation of the diffusion-based approach for super-accurate local alignments.

Following your observation, we do more analysis using the results from Table 1. As presented in the table below, the improvement margin for rotation is relatively modest (+9.93%), whereas other metrics show significant gains exceeding 20%.

In addition, the failure cases illustrated in Figure 7 are also highly related to the high rotation error:

1. Local geometric ambiguity: Similar geometry across fragments makes it difficult to determine precise rotations.
2. Small Fracture Surfaces: Tiny pieces often lack distinct surface features, leading to 180-degree rotation errors.

All these results demonstrate that our diffusion-based method does not handle rotation as effectively as traditional optimization-based approaches. We provide some thoughts below:

The other three metrics heavily rely on an accurate translation/placement of fragment pieces, where the diffusion-based approach excels by learning global shape priors together with local alignments. On the contrary, the RMSE of rotation mainly examines the accuracy of fine-grained alignments, which relies more on accurate local shape matching. It seems that methods like Jigsaw can produce better rotation-level alignments by conducting direct optimization based on the local surface matching results, while our diffusion-based approach allocates most of the learning capacity for global shapes (correct translations).

We believe that a potential direction for improving the rotation-level accuracy can be adding an additional stage that focuses on refining the rotation parameters. We will include this discussion in the paper.

W1 & Q3: Why does PuzzleFusion++ significantly outperform DiffAssemble, even among diffusion models? Can the authors describe the key difference that make your model works better?

We presents our core design of denoiser in the ablation section (L454-L466), which investigates 3 core modules of our SE3 denoiser. 1) We have the pre-traiend PointNet++ VQ-VAE to encode local geometric information, while DiffAssemble only encodes a global semantic latent for each fragment, thus lacking fine-grained local details. 2) DiffAssemble does not take care of the ambiguity of 3D rigid transformation, while we use the anchor fragment to avoid the training ambiguity. 3) We have a tailored noise scheduler to handle the 3D fracture assembly task while DiffAssemble simply keeps the default designs of DDPM.

We appreciate your feedback and suggestion. We have included the rotation discussion in Appendix B.4 in the updated PDF version.

We thank you for the construction comments and the questions.

W1: The Figure 3 is a little puzzled. A better format is recommended.

We thank for the suggestion. Reviewer ox6f also mentioned making Figure 3 more concise. We have updated pdf with new version Figure 3. We reformat figure 3 to the inference pipeline instead of seperate components.

W2&Q6: As mentioned in Weakness, could you provide more complicate objects during rebuttal?

We thank for the suggestion. Artifact is the other subset in the Breaking Bad dataset which contains many complicated objects than the Everyday subset. To have a proper evaluation on the Artifact subset, we take the model pretrained on the Everyday subset, and finetune it on the training split of the Artifact subset with only 20% of the pretraining iteration, and then evaluate the model on the test split of Artifact subset.

The qualitative and quantitative results are provided in Figure 17 and Table 9, showcasing the model's performance on these more challenging objects

Q1: The main difference between PuzzleFusion and PuzzleFusion++ should be clarified in the paper, since it is a future work from PuzzleFusion.

PuzzleFusion is the first work that employs diffusion models to solve 2D spatial puzzles, demonstrating the potentials of diffusion models as general iterative solvers for non-generation or discriminative tasks. Our paper further extends diffusion models to more complicated 3D fracture assembly problem.

At architecture level, PuzzleFusion is specifically designed for 2D jigsaw puzzles with simple polygonal shapes, where a 1D chain of 2D coordinates represents each puzzle piece. The approach cannot be used for the 3D shape assembly task, where each piece is a set of unordered 3D points.

In addition, the 3D fracture assembly task poses several challenges: i) The 6-DoF solution space is more complicated than the 3-DoF solution space of 2D spatial puzzles; ii) Unlike 2D puzzle solvers that mainly leverage image semantics or polygon structures, 3D fracture assembly requires a deep understanding of fracture surfaces; iii) There can be many small 3D fragments, further complicating the problem. Our PuzzleFusion++ proposes new designs to address these challenges. More specifically, our VQ-VAE with PointNet++ can encode fine details of fracture surfaces into a set of local latents, boosting the shape understanding of the denoising network; Our auto-agglomerative framework verifies assembly results and composes confident small pieces into larger ones to facilitate future iterations, which effectively improves the assembly success rate; We also design tailored components for the diffusion model, including the noise scheduler and the denoising transformer with local shape encodings as condition.

Q2: Are 25 points enough to represent a fragment in Sec.3.1?

Following the experimental setting in the literature, each fragment piece has 1000 sampled points, and our PointNet++ VQ-VAE encode the 1000 points into the 25-point latent vector. Each point here has a latent embedding that captures local shape details. One way to better understand whether the 25-point latent representation is to examine the quality of the reconstructed point cloud. As shown in Figure 8, the visualization demonstrates that the decoded point clouds retain the overall shape and essential details of the original fragments.

Q3: Does the pairwise alignment verifier work well for all objects?

We have included the pairwise alignment verifier's performance ablation in Table 6. The verifier does not work well for all objects. To handle this, we only classified scores > 0.9 as the true label, achieving 90.77% accuracy, 87.88% precision, and 45.95% recall. This high threshold ensures that the selected pairs have high confidence, lowering the possibility of wrong merging.

Additionally, Table 5 shows the upper-bound performance of our method using the ground truth verifier. The ground truth verifier performs significantly better than the verifier using Jigsaw matchings, highlighting the potential for further improving the verifier.

Q4: In auto-agglomerative inference, are 6 iterations enough for merging? Have you try more iterations?

Yes, the 6 iterations are enough for merging. We tested with more iterations and found that increasing the iterations did not improve performance.

Below are the results of more auto-agglomerative iterations:

Q5: Please discuss could reinforcement learning whether is help for this task.

We believe RL is less suitable for this task due to the challenges of defining an effective reward function. The final reward can only be provided after completing all assembly steps, making it hard to optimize the learning process. This limitation could slow down learning and reduce the overall system performance. We will include this discussion in our final version.

Q7: Please discuss the possible solution for your limitations.

Local Geometric Ambiguity: The issue of local geometric ambiguity arises because some fracture surfaces are too similar and “confuse” the model. A possible solution is to increase the number of sampled 3D points for each fracture (currently we use 1000, the default setting inherited from previous works in the literature). By increasing the point cloud density, the PointNet++ encoder can capture more accurate local geometric details, and the denoising network can have better chance to distinguish these similar fracture surfaces.

Small Fracture Surfaces: Small fracture surfaces can result in failed connections between merged fragments. A potential solution is to enable the network to learn more global information. If the network can imagine the shape of original object, it can correctly place the fragments in their intended positions.

We will include this discussion into the future work.

Thank you for the clarifications and for addressing my concerns. I will keep my rating as positive.

We appreciate the reviewer’s constructive and detailed feedback. We try to address the concerns and questions below.

Strengths 2: More discussion/analysis on the noise scheduler.

We thank you for the insightful comment. We agree that our analysis of the scheduler was limited. We have added quantitative and qualitative comparisons of three schedulers (i.e., linear, cosine, and ours) together with the visualizations of the denoising process in Table 8 and Figure 11 and Figure 10.

Our noise scheduler is used for both training and testing, as we follow the vanilla DDPM formulation rather than EDM[A]. With our tailored noise scheduler for the 3D shape assembly task, the denoising process allocates more steps to refining precise local adjustments rather than finding the rough global location (at test time). As for training, this scheduler indeed makes the process more efficient, as fewer training iterations are spent on "less important" timesteps. If the default (linear or cosine) scheduler were used for training while our scheduler was applied during testing, similar results might still be achieved but would require more training iterations.

To illustrate the advantages of our scheduler during the testing stage, we compare it to the linear and cosine scheduler. The linear scheduler uses most of its steps (T=1000 to 150) for rough localization (Figure 10), while the cosine scheduler allocates more denoising steps in the final adjustment phase, outperforming the linear scheduler by a clear margin (Table 8).

Building on this idea, our scheduler dedicates an even larger portion of the denoising steps to the final adjustment, further improving results (Table 8). The top three rows of Figure 11 show simpler cases of objects comprising at most 5 fragments. While all the schedulers achieve 100% part accuracy, gaps between fragments are visible for the linear or the cosine schedulers. Our precise alignments may have minimal effects on the standard metrics but significantly enhance the quality of the final assembly.

W1: Fair comparison regarding anchor fragments

We believe the reviewer's main concern is that the post-processing step introduced in Jigsaw could lead to unfair evaluations. This step involves using the largest fragment to align the predicted assembly with the ground truth before metric calculation. Both our method and Jigsaw employ this post-processing. We also noticed that the recent ICML'24 paper "3D Geometric Shape Assembly via Efficient Point Cloud Matching", as mentioned by Reviewer ox6f, also uses this setting. This is the "proper" test setting -- those early baseline methods should have done it correctly.

To correct the testing setting of those baselines, we applied the same post-processing step to all of them (See the table below). It is evident that our method still significantly outperforms the baselines. We also included the numbers reported in the original submission in brackets.

Regarding whether a method uses the anchor fragment during training, we think it's the design choice of each method and has nothing to do with fairness. We added this design component to handle the ambiguity of rigid 3D transformation for our method, and we are not responsible for adding this to all baselines. Fairness should be defined by whether the task input and output are the same across methods. Our input consists of fragments’ point clouds with random poses, and the output is the final alignment parameters for each point cloud, which are identical to those of the baselines. We do not require additional information about the anchor fragment or its ground truth pose beforehand.

W2: Missing baselines

We apologize for the oversight. We checked the two papers and discussed them here and will include them in our final version.

• PHFormer:

  The AAAI paper introduces PHFormer, which employs a hybrid attention module to model the relationships between fragments. We compared our performance with theirs in the table below:

  Method	RMSE (Rot.) ↓	RMSE (Trans.) ↓	PA ↑	CD ↓
  PHFormer	26.1	9.3	50.7	9.6
  Ours	38.1	8.04	70.6	6.02

  Our method outperforms PHFormer in three metrics but has a higher RMSE in rotation. We will also include these quantitative and additional qualitative comparisons with PHFormer in the final version.

• Fracture Assembly with Segmentation And Iterative Registration:

  The ICASSP paper introduces FRASIER, a framework for reassembling fractured objects using fracture surface segmentation and iterative registration. For the registration stage, FRASIER uses the GeoTransformer[B] to match points between fragment pairs.

  Unlike prior methods that use 1k points per fragment, FRASIER samples 50k points per fragment. This high point density enables detailed fracture surface capture but does not align with the standard 1k-point setup used by other methods. Moreover, since their code is not available, we cannot evaluate their performance under the 1k-point setting.

  To infer FRASIER’s performance with 1k points per fragment, we refer to GeoTransformer’s reported performance in Jigsaw (Appendix E.3). With 1k points, GeoTransformer produces poor results (RMSE (Rot.) = 84.8°, RMSE (Trans.) = 14.3 × 10⁻², PA = 3.1%). These findings show that GeoTransformer is ineffective under the 1k-point setting. Since FRASIER relies on GeoTransformer for registration, it likely cannot achieve reasonable results under the 1k-point-per-fragment setting.

Q1: What guarantees that the input representation are equivariant?

The input representation is equivariant because of the way we pre-trained our VQ-VAE. During pretraining, we applied random rotations to the input point cloud and used the rotated point cloud for supervision. This ensures that the latent embedding reflects the rotation applied to the input, achieving equivariance empirically.

[A] Karras, T., Aittala, M., Aila, T., & Laine, S. (2022). Elucidating the design space of diffusion-based generative models. Advances in neural information processing systems, 35, 26565-26577.

[B] Qin, Z., Yu, H., Wang, C., Guo, Y., Peng, Y., Ilic, S., ... & Xu, K. (2023). Geotransformer: Fast and robust point cloud registration with geometric transformer. IEEE Transactions on Pattern Analysis and Machine Intelligence, 45(8), 9806-9821.

We thank the reviewer for the constructive and detailed feedback. We provide the response to each point below.

W1: Providing Figure 3 as a form of an algorithm may improve its readability.

We thank for the suggestion. We are working on redesigning Figure 3, and we first address the remaining questions/concerns in this thread. We will update Figure 3 in the pdf shortly and send a follow-up message.

W2: Overstatement to mention that anchor initialization 'does not affect' the quality of the assembly results

We apologize for the confusion — Table 7 only shows results with #iteration=1, so it actually supports our claim about the robustness over different anchor initializations. To clarify, we add results for #iteration=2 and #iteration=4 with random anchor initializations. Similar to the setup of Table 7, we run the inference 10 times and calculate the mean and variance. For better readability, we also included the numbers reported in Table 2 of the main paper in brackets. This confirms the quality consistency across different initializations.

W3: Missing baseline

We thank you for pointing out the ICML paper. We apologize for the oversight.

The ICML paper presents a new method, PMTR, which employs an efficient high-order feature transformation layer to establish reliable correspondences. When investigating the paper and the official implementation, PMTR used an easier "volume-constrained" subset of the Breaking Bad dataset, where fragments below a minimum volume threshold are excluded. Most previous works, including our submission, use the original, more difficult dataset. To make a fair comparison, we train and evaluate our method using the volume-constrained subset and provide the results below:

PuzzleFusion++ outperforms PMTR on all metrics except rotation error (RMSE(R)) by a small margin (approximately 1 degree). Please refer to Appendix B.3 for more details.

W4: Typos

We thank you for pointing out these. We fixed the typos in the updated PDF.

Q1: Optimal performance of the number of iterations.

The maximum iteration was set to 6 as it is close to convergence, with minimal or no consistent improvements further. Below are the results for more iterations.

These results confirm that iteration 6 achieves nearly optimal performance, with further iterations yielding negligible or inconsistent changes.

Q2: What are the results when the scheduler is x

The ablation of the noise scheduler had already been included in the first row of Table 4 in the initial submission. When the scheduler is "×," we use a linear scheduler by default.

Q3: Could the authors include a quantitative/qualitative comparison of all 3 noise schedulers?

We appreciate the suggestion. We updated the PDF by adding a quantitative comparison of the linear, the cosine, and our schedulers in Table 8 and a qualitative comparison in Figure 11.

Additionally, we include visualizations of the denoising process with different schedulers in Figure 10. These visualizations further illustrate how the linear and cosine schedulers spend more time finding the rough locations of fracture pieces, while ours focus more on precise alignment.

Thank you for your detailed response! I believe my concerns and questions have been addressed adequately. I have raised my recommendation.