SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Rigidity Loss. As 3D fracture assembly is applied to rigid objects, we further enforce the rigidity loss over the pairs of matched pieces. Let  $\tilde{X}_{ij} \in \mathbb{R}^{\hat{N}_i \times \hat{N}_j}$  be the submatrix of  $\hat{X}$  that indicates the likelihood between the fracture points on the piece  $P_i$  and  $P_j$ . For a fracture point  $p$  in  $P_i$ ,  $\tilde{X}_{ij}$  matches it to a point  $p'$  computed as the weighted average of all the matchable points in  $P_j$ :
$$
p ^ {\prime} = \frac {\tilde {X} _ {i j} (p) P _ {j}}{\| \tilde {X} _ {i j} (p) \| _ {1}}, \tag {9}
$$
where  $\tilde{X}_{ij}(p) \in \mathbb{R}^{1 \times \hat{N}_j}$  is the row in  $\tilde{X}_{ij}$  that contains matching likelihood from  $p$  to  $P_j$ . We optimize the transformation  $\tilde{T}_{ij}$  from  $P_i$  to  $P_j$  by minimizing the weighted mean squared error between the matched fracture points
$$
\tilde {R} _ {i j}, \tilde {\boldsymbol {t}} _ {i j} = \underset {R, \boldsymbol {t}} {\operatorname {a r g m i n}} \sum_ {p \in P _ {i}} \| \tilde {X} _ {i j} (p) \| _ {1} \| R (p) + \boldsymbol {t} - p ^ {\prime} \| _ {2}. \tag {10}
$$
where  $\tilde{R}_{ij},\tilde{t}_{ij}$  are the rotation and translation part of  $\tilde{T}_{ij}$  and  $R(\cdot)$  is the rotation operator. (10) has a closed-form solution:
$$
(U, \Sigma , V ^ {\top}) = \operatorname {S V D} \left(\sum_ {p \in P _ {i}} \| \tilde {X} _ {i j} (p) \| p ^ {\top} p ^ {\prime}\right), \tilde {R} _ {i j} = U V ^ {\top}, \tilde {\boldsymbol {t}} _ {i j} = - \frac {\sum_ {p \in P _ {i}} \| \tilde {X} _ {i j} (p) \| _ {1} \left(R (p) - p ^ {\prime}\right)}{\sum_ {p \in P _ {i}} \| \tilde {X} _ {i j} (p) \| _ {1}}. \tag {11}
$$
The rigidity loss is computed as:
$$
\mathcal {L} _ {\mathrm {r i g}} = \sum_ {1 \leq i, j \leq n} \mathcal {R} _ {i j}, \text {w h e r e} \mathcal {R} _ {i j} = \sum_ {p \in P _ {i}} \| \tilde {X} _ {i j} (p) \| _ {1} \| \tilde {R} _ {i j} (p) + \tilde {t} _ {i j} - p ^ {\prime} \| _ {2}. \tag {12}
$$
To avoid numerical instability and improve efficiency in (11) during training, we compute  $\tilde{T}_{ij}$  in each iteration without involving back-propagation and only set  $\tilde{X}_{ij}$  as the learnable variable.
The loss for training is composed as:
$$
\mathcal {L} = \alpha \mathcal {L} _ {\text {s e g}} + \beta \mathcal {L} _ {\text {m a t}} + \gamma \mathcal {L} _ {\text {r i g}}. \tag {13}
$$
Multi-part Matching With the doubly-stochastic matrix  $\tilde{X}$ , computing multi-part matching results becomes straightforward. Given the fact that the primal and dual feature represents the surface looked from a different viewpoint, they should be distinct and not assigned good confidence for matching. Therefore, we can directly pass  $\tilde{X}$  to a Hungarian layer [37]. This yields a binary permutation matrix  $X$ , representing a bipartite matching among all fracture points:
$$
X = \operatorname {H u n g a r i a n} (\tilde {X}) \tag {14}
$$
As mentioned, the resulting matching matrix  $X$ , along with the affinity matrix  $M$  and soft matching matrix  $\tilde{X}$ , captures the local geometry similarity across different pieces. Additionally, they quantify the confidence level of how well the pieces fit together. This information enables an accurate positioning and alignment of the fractured components.
3.4 Global Fracture Alignment
To efficiently and robustly recover the global poses of each fractured piece of object  $O$ , we adopt a two-step pipeline based on the bipartite matching  $X \in \{0,1\}^{\hat{N} \times \hat{N}}$  predicted by the multipart matching module.
In the first step, we compute pairwise transformations between each pair of pieces  $(P_{i}, P_{j})$  to remove outlier matches. Let  $X_{ij} \in \mathbb{R}^{\hat{N}_{i} \times \hat{N}_{j}}$  denote the submatrix of  $X$ , which establishes the correspondences between the fracture points from  $P_{i}$  and  $P_{j}$ . We apply the RANSAC algorithm [38] to compute the transformation  $\tilde{T}_{ij}$  from  $\hat{P}_{i}, \hat{P}_{j}$  and  $X_{ij}$ . This step ensures reliable pairwise transformations for subsequent alignment.
In the second step, we perform robust global alignment using the computed pairwise transformations. We model the global alignment configuration as a factor graph [39], denoted as  $G = (\mathcal{V}, \mathcal{E})$ , where the global poses  $\tilde{T}_i$  are optimized on the vertices  $\mathcal{V}$  and the pairwise transformations  $\tilde{T}_{ij}$  are set as constraints on the edges  $\mathcal{E}$ . The information matrix  $I(e)$  over the edge  $e = (i,j)$  is set to  $|X_{ij}|_{\mathrm{Fro}}^{-2} \cdot I_6$  with  $I_6$  being the  $6 \times 6$  identity matrix. We employ Shonan averaging [40], a state-of-the-art 6-DoF global alignment method, to optimize global poses over  $G$ . The global alignment method takes pairwise transformations between potentially adjacent pieces as input and outputs a pose for each piece up to a global transformation over the coordinate system. To avoid such an ambiguity, we anchor the coordinate system to the canonical one of the largest pieces in our experiment and measure the errors over all other pieces of the fractured object. Further technical details of this module can be found in the references [41, 42, 43]. Please refer to those sources for more information, as it is beyond the scope of our main contribution.
We demonstrate the effectiveness of our 3D fracture assembly framework through experimental evaluations on a large-scale fracture assembly dataset. Our results show that Jigsaw significantly outperforms the baseline methods both quantitatively and qualitatively. Additionally, we conduct ablation studies to analyze the contributions of each module in our framework. All experiments are conducted on a Linux workstation with 4 Tesla V100-SXM2-32GB GPUs, Intel(R) Xeon(R) CPU E5-2698 v4 @ 2.20GHz CPUs, and 480GB Memory. Table 1 lists the configurations of the parameters in our experiments. A comparison of the training/testing time is included in Appendix D.
4.1 Protocols
Dataset. We leverage the Breaking Bad dataset (Sellan et al., 2022), a novel data set of multiple fracture assemblies featuring synthetic physical breaking patterns. Our training was on everyday subset and the testing was on both everyday and artifact subsets for a fair comparison with
Table 1: The detailed experiment parameters. We follow the parameters provided in [16, 33, 4] to reproduce their results. Training, model, and dataset parameters are included.
Jigsaw	Predator[16]	DGL [33]	LSTM [33]	Global [33]	description
epoch	250	200	200	200	200	training epochs
bs	4	16	32	32	32	batch size
lr	0.001	0.01	0.001	0.001	0.001	learning rate
optimizer	Adam	SGD	Adam	Adam	Adam	optimizer during training
schedule	Cosine	Exponential	Cosine	Cosine	Cosine	learning rate scheduler
min_lr	1e-5	-	1e-5	1e-5	1e-5	minimum learning rate for Cosine scheduler
α	1.0	-	-	-	-	segmentation loss ratio
β / e	1.0 / 10	-	-	-	-	matching loss ratio change at epoch e
γ / e	1.0 / 200	-	-	-	-	rigidity loss ratio change at epoch e
τ	0.05	-	-	-	-	temperature parameter for affinity
sampling by points (N)	object	piece	piece	piece	piece	sampling strategies for point clouds
5000/o	800/p	1000/p	1000/p	1000/p	points sampled per object (/o) or piece (/p)
η	0.025	0.025	-	-	-	segmentation ground truth label threshold
bases. Everyday consists of 498 models and 41,754 fracture patterns, and is split into a training set (34,075 fracture patterns from 407 objects) and a test set (7,679 fracture patterns from 91 objects). Artifact consists of 3651 fracture patterns from 40 un categorized objects. The average diameter of the objects in the training and testing dataset was 0.8. Categorical information was concealed during all the experiments.
The input was uniformly sampled  $N$  points on the surface of the object as a point cloud. In composing the point cloud, two sampling strategies are considered, sampling by piece and object. Sampling by piece, as used in the Breaking Bad benchmark [4], equally samples points in each fragment, which would cause point density imbalance. To better reflect real-world scanning, we opt for "sampling by object", sampling a fixed number of points in each object, while sampling points for each fragment based on its surface area. We ensure at least 30 points per fragment for multi-part matching. Additional processing details can be found in the Appendix C.1.
Baseline Approaches. As learning to assemble multiple fractured objects is a novel task, there is a dearth of existing methods for direct comparison. Therefore, we have selected a set of established methods that primarily address similar challenges in 3D alignment and assembly, and employ them as baselines for comparative evaluation. Global [26, 10] extracts per-piece features, which are combined with global shape descriptors to regress the pose of each piece in one shot. LSTM applies a bidirectional LSTM similar to [27] and estimates the pose of each piece in sequential style. DGL [8] is a state-of-the-art approach for the assembly task of parts and can also be adapted to the assembly task of fractures. It leverages an iterative graph neural network to reason about the relationships among pieces. Additionally, we take PREDATOR [16] as a minor competitor, which is a state-of-the-art approach for 3D registration with low overlapping rates. It employs a multitask transformer to estimate the overlap between pairs of pieces in conjunction with their relative pose. All baseline approaches use textureless point clouds as input and are trained under the everyday object subset of the Breaking Bad dataset.
Evaluation Metrics. We adopt the same evaluation schemes for 3D assembly tasks used in [4, 14, 10]. We report the mean absolute error (MAE) and the root mean square error (RMSE) for the rotations and translations of the estimated global poses. Additionally, we include the part accuracy (PA) metric proposed in [33], which measures the ratio of perfectly assembled pieces based on an average Chamfer distance of less than 0.01 for each point between the assembly results and the ground truth.
4.2 Performance
Overall Performance. We report the performance of Jigsaw and all the baseline methods over the 3D fracture assembly task on both everyday and artifact object subsets. An overall quantitative comparison of the evaluation metrics is presented in Table 2. For PREDATOR we apply the evaluation metrics on pairwise transformations since it only supports overlapped pairs of pieces as input.
Jigsaw significantly outperformed all the baselines across all the evaluation metrics. On everyday, we achieved an average rotation error of  $36.3^{\circ}$  over all pieces, which was a  $46\%$  reduction compared to the top-performing baseline method, DGL. Our transformation error was also outstanding, with an average of  $8.7 \times 10^{-2}$ , surpassing the best baseline method with an error of  $11.8 \times 10^{-3}$ . We
Table 2: Quantitative results of baseline methods and Jigsaw on the Breaking Bad dataset.
Method	Original Task	RMSE (R) ↓ degree	MAE (R) ↓ degree	RMSE(T) ↓ ×10-2	MAE (T) ↓ ×10-2	PA ↑ %
Results on the everyday object subset.
Global [26, 10]	assembly	82.4	69.7	14.8	11.8	21.8
LSTM [27]	assembly	84.7	72.7	16.2	12.7	19.4
DGL [8]	assembly	80.6	67.8	15.8	12.5	23.9
DGL(5000/o) [8]	assembly	81.1	68.1	15.4	12.3	25.5
PREDATOR [16]	registration	82.8	71.2	12.5	10.2	1.3
Jigsaw (Ours)	assembly	42.3	36.3	10.7	8.7	57.3
Results on the artifact object subset.
Global [26, 10]	assembly	86.9	75.3	17.5	14.5	5.6
LSTM [27]	assembly	85.6	74.1	18.6	15.2	4.5
DGL [8]	assembly	86.3	74.3	18.0	14.9	9.6
PREDATOR [16]	registration	86.0	74.8	13.4	10.9	1.1
Jigsaw (Ours)	assembly	52.4	45.4	22.2	19.3	45.6
Figure 3: Qualitative results of baseline methods and Jigsaw on the Breaking Bad dataset (mesh used only for visualization). The coordinate system of the green piece is set as the global coordinate system. Better viewing with color and zooming in.
successfully restored over  $57\%$  of the pieces to their original pose while baseline methods restored only half of us. We also show that change in sampling strategy has had no impact on the outcome of baseline methods through a sampling by object version of DGL. Our method exhibited remarkable generalizability on artifact, showing a slight increase of  $10.1^{\circ}$  in average rotation error and achieving a part accuracy of  $44\%$ . In contrast, all baseline methods completely failed to handle new categories of models that were not present in the training dataset. The PREDATOR model failed to accurately identify the pairwise overlap region, resulting in erroneous predictions of complete overlap between the two pieces. While this may seem to yield a small translation error, it failed to perform adequately in the assembly scenario. Hence, we exclude it from further discussion.
Detailed distribution of each metric. Fig. 4 presents a comprehensive analysis of each metric for all assembly models. Regarding rotation and translation metrics, a larger surface area under the curve indicates better performance. Jigsaw exhibits a significantly larger number of samples with small rotation and translation errors, as illustrated in the figure. Although Jigsaw reports a slightly larger translation error than the baseline models on the artifact object subset, the distribution results reveal that this is primarily due to a few extreme cases, leading to an undesirable accumulation of the translation error. In terms of part accuracy, a larger surface area above the curve signifies superior results. Notably, Jigsaw demonstrates a higher proportion of correctly positioned pieces.
Visualization. The qualitative comparison of different models is shown in Fig. 3. We plot the recovered object from different models. It demonstrates that our model could find the position of the fracture pieces and recover a good pose of each piece. Visualization of objects from different categories also implies our model has good generalization ability and it doesn't require object-level information as assembly guidance.
Figure 4: The discretized distributions of metrics evaluated over the test data of everyday (up) and artifact (bottom) object subset in the Breaking Bad dataset.
Ablation Study. We conducted an ablation study to assess the effectiveness of each component in our Jigsaw model. The baseline model consists of a segmentation module and a multi-part matching module, trained separately using separate backbones and without incorporating transformers. The plain joint training model was trained using the same loss function as the Jigsaw model, but without using two transformer layers. The effectiveness of the rigidity loss regularization was evaluated over the complete network with both joint training and transformer layers.
As presented in Table 3, the joint training model outperformed the baseline in transformation metrics. This is due to a more favorable initialization for the matching module achieved through pretraining the backbone with the segmentation module. Joint training also aided the matching module's convergence to better optimum. Attention layers significantly reduced rotation errors:
the self-attention layer embedded geometric information into point features, while the cross-attention layer improved multi-part matching by incorporating information from other parts. Applying rigidity loss as a refinement additionally brought an improvement of  $2 \times 10^{-3}$  in translation. Even the baseline method outperformed previous works, highlighting the effectiveness of our segmentation and multi-part matching modules.
Table 3: Ablation study of Jigsaw.
Components	RMSE (R)	MAE (R)	RMSE (T)	MAE (T)	PA
Joint Training	Attention	Rigidity	degree	degree	×10-2	×10-2	%
51.9	44.6	32.4	27.2	45.7
✓	51.7	44.4	12.7	10.2	47.0
✓	✓	42.3	36.3	10.9	8.9	57.2
✓	✓	✓	42.3	36.3	10.7	8.7	57.3
In this study, we propose a novel geometry-aware framework to address the 3D fracture assembly task. Our approach involves a joint learning model that effectively extracts local geometry features from the backbone and enables the simultaneous learning of piecewise fracture surface segmentation and global fracture point matching. Experimental evaluations and thorough analyses substantiate the exceptional performance of our method in terms of accurately recovering piece poses and restoring the original object. Furthermore, we demonstrate the generalizability of our approach across objects of diverse categories. One limitation of our approach lies in that as the number of fractured pieces increases, the accuracy of pose estimation tends to decrease due to potential ambiguities in fracture matching caused by local geometric features. Future research directions may include exploring matching pruning techniques for pairwise pose estimation to prevent penetration between fractured pieces and enforcing smoothness constraints to enhance pose accuracy.
We would like to acknowledge NSF IIS-2047677, HDR-1934932, and CCF-2019844.
[39] Daphne Koller and Nir Friedman. Probabilistic graphical models: principles and techniques. MIT press, 2009.
[40] Frank Dellaert, David M Rosen, Jing Wu, Robert Mahony, and Luca Carlone. Shonan rotation averaging: Global optimality by surfing  $so(n)$ . In European Conference on Computer Vision, pages 292-308. Springer, 2020.
[41] Qian-Yi Zhou, Jaesik Park, and Vladlen Koltun. Open3d: A modern library for 3d data processing. arXiv preprint arXiv:1801.09847, 2018.
[42] Frank Dellaert and GTSAM Contributors. borglab/gtsam, May 2022.
[43] Frank Dellaert and Michael Kaess. Factor Graphs for Robot Perception. Foundations and Trends in Robotics, Vol. 6, 2017.
A Details of Network Architecture
We provide additional information about our network architecture. The backbone of our network is based on a standard PyTorch implementation of multi-scale grouping PointNet++. It is followed by a single MLP layer that extracts a  $D$ -dimensional feature for each point. In our implementation, we set  $D = 128$ .
The detailed structure of the transformer module can be seen in Fig. 5. The self-attention layer follows the official code of the point transformer layer, while the cross-attention layer employs a multi-head attention mechanism with position-wise feed-forward networks. In our configuration, we set the number of attention heads to  $h = 8$ , the head dimension to  $d_h = 16$ , and use  $k$ -nearest neighbor sampling with  $k = 16$ . The inner layer of the feed-forward network has a dimensionality of  $d_i = 256$ .
Figure 5: Detailed structure of the attention layers.
Regarding the primal-dual descriptor, we set its feature dimension to  $d = 256$ .
B Analysis on Primal-dual Descriptor
Figure 6: Comparison between the doubly-stochastic soft matching matrix  $\tilde{X}$  computed by different types of descriptors.
We conducted an analysis to compare the results obtained using the primal-dual descriptor and the standard single-view descriptor. Before passing the affinity metric computed from the single-view descriptor to the Sinkhorn layer, a common practice is to mask out the diagonal sub-matrix to prevent self-self alignment. For clarity, we refer to the masked and unmasked versions of the primal-dual and single-view descriptors to denote whether this masking operation has been applied or not.
In Fig. 6, we provide a visualization of the doubly-stochastic matrix  $\tilde{X}$  computed by the Sinkhorn algorithm using four types of descriptors: (1) the unmasked primal-dual descriptor, (2) the masked primal-dual descriptor, (3) the unmasked single-view descriptor, and (4) the masked single-view descriptor. The visualization depicts an example object with 3 pieces and 58 fracture points. The  $\tilde{X}$  obtained from the primal-dual descriptor clearly demonstrates its ability to differentiate between the two viewpoints and avoid aligning a point to itself. In contrast, the single-view descriptor exhibits a diagonal peak, resulting in self-self alignment. Even when the diagonal sub-matrix is masked to prevent self-self alignment, the soft matching computed from the single-view descriptor is less distinct compared to that computed from the primal-dual descriptor.
Furthermore, we conducted evaluation on the everyday object subset of Breaking Bad dataset, using the unmasked primal-dual descriptor and the masked primal-dual descriptor. The results indicate that the unmasked primal-dual descriptor achieves comparable performance to the masked one, with only a slight difference: a decrease of  $0.2^{\circ}$  in rotation error (MAE(R)) and an improvement of  $0.7 \times 10^{-2}$  in translation error (MAE(T)). We attribute this minor difference to the fine characteristics of the primal-dual descriptor that prevent self-self alignment, thus resulting in consistent performance.
We leverage Breaking Bad dataset [4] to evaluate our method and all baseline methods. As we have stated in section 4, the training of all methods were on the training subset of everyday, and testing were on the testing subset of both everyday and artifact object. Each object within the dataset has been fragmented into pieces, represented by triangle meshes, and all the pieces are in their original poses. By assembling these pieces together directly, the surface of the original object is seamlessly restored. The triangle meshes of the pieces solely consist of exterior faces that are visible from the outside.
In generating the point cloud for our experiments, we employ two sampling strategies for the compared methods: "sampling by piece" and "sampling by object". The "sampling by piece" strategy, originally utilized in the Breaking Bad benchmark [4], involves sampling an equal number of points within each fragment. However, this approach leads to excessively dense sampling on small fragments, while larger fragments suffer from sparser point distributions, resulting in an imbalance in the representation of point density across fragments. To better mirror real-world scanning technology, we opt for the "sampling by object" strategy. With this approach, we sample a fixed number of points within each object, and the number of sampled points for each fragment is determined based on its surface area. In essence, smaller fragments receive fewer points, whereas larger ones receive more, ensuring a more realistic representation. Additionally, we ensure that each fragment is sampled with a minimum of 30 points to include even the tiniest fragments in multi-part matching. Detailed parameter settings can be found in Table 1. Our analysis of the average fragment numbers and experiments conducted using DGL [8] indicate that the choice between the two sampling strategies has negligible impact on the baseline performance.
For each sampled point  $p$  on  $P_{i}$ , its fracture label  $c_{p}$  is determined by the distance from  $p$  to its nearest neighbor  $q$  among points from all other pieces:
$$
c _ {p} = \mathbb {1} \left(\min  _ {q \in O \backslash P _ {i}} \| p - q \| _ {2} <   \eta\right) \tag {15}
$$
where  $\eta$  is set to 0.02 for all the objects. The ground-truth matching point  $\hat{q}$  of a fracture point  $\hat{p} \in \hat{P}_i$  is set to
$$
\hat {q} = \underset {\hat {q} ^ {\prime} \in \hat {O} \backslash \hat {P} _ {i}} {\operatorname {a r g m i n}} \| \hat {p} - \hat {q} ^ {\prime} \| _ {2}. \tag {16}
$$
where  $\hat{O},\hat{P}_i$  denotes the set of fracture points in  $O$  and  $P_{i}$  During the training process,  $c_{p}$  was applied to segment the fracture points and during testing the predicted  $\tilde{c}_p$  was used instead. An example is shown in Fig. 7.
After the ground-truth labeling and matching were com
puted based on the pieces at their original pose, all pieces were recentered to the origin and a random rotation was applied to each piece.
To ensure a fair comparison with baseline methods [26, 10, 27, 8], we adopted the same implementation as of the benchmark code of [4], which sampled the same number of point each piece. For PREDATOR, we applied the same sample routine as the baseline methods, and we paired the adjacent pieces for training and testing.
C.2 Evaluation Metrics
The mean absolute error MAE(R) and square-rooted mean squared error RMSE(R) of rotation were computed as