SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Harel et al. Harel and Ben-Shahar (2021) and Shabani et al. Shabani et al. (2021) are state-of-the-art methods for CJP and RLA, respectively. We have used their public implementations to compare in the corresponding tasks. Note that Harel et al. is not applicable to VJP or RLA, where neighboring angles may not add to  $180^{\circ}$  or corners may not meet. Shabani et al. is not applicable to CJP or BJP due to its poor scalability (i.e., exponential in the number of pieces). We have also prepared a third method based on the transformer network to be compared for RLA. The following provides more information, while the full details are in supplementary.
Harel et al. Harel and Ben-Shahar (2021) proposed a two-step algorithm for CJP: Enumerating pairs of compatible pieces by heuristics based on edge lengths or corner angles, then globally solving for the whole arrangement by a spring system.
Shabani et al. Shabani et al. (2021) enumerates arrangement candidates by heuristics and learns to regress the realism of an arrangement candidate by deep neural networks. Since the number of room/door types differs in our datasets, we made minor modifications to the data loader and the network architecture.
TransVector is a baseline Transformer network with a vector representation that directly estimates the pose parameters instead of iterative denoising. TransVector shares the same architecture as our denoising network at the core with the following changes: 1) Remove time-dependent features  $(x_{i,t}^{r})$  and  $t$ ) from the embedding (3); 2) Change the supervision  $\delta_{i,t}^{r}$  from the noise to the position/rotation parameters; and We have also compared with the Transformer network with a raster representation, which is presented in the supplementary.
5.3 Main results
Table 1 compares the proposed approach against Shabani et al. Shabani et al. (2021) and TransVector. Shabani et al. runs exponentially in the number of pieces (i.e., rooms), taking hours or even days to process a single house with seven or more rooms. Therefore, we collect small houses (i.e., at most six rooms) to create "Small RPLAN" and "Small MagicPlan" datasets for its evaluations. For each experimental setting (e.g., Small MagicPlan), we train a network for each method.
Table 1: RLA quantitative results with two metrics: Positional Error (MPE) and Graph Editing Distance (GED). Small RPLAN (resp. Small MagicPlan) is a subset of the corresponding full dataset, consisting of houses with at most 6 rooms. The small datasets are created for Shabani et al. Shabani et al. (2021), which is not scalable to many rooms.
Dataset	Small RPLAN	Full RPLAN	Small MagicPlan	Full MagicPlan
Metric	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)
Shabani et al.	29.44	1.28	✗	✗	36.63	1.89	✗	✗
TransVector	36.09	1.51	46.18	2.27	40.80	2.38	53.11	6.41
Ours	8.65	0.90	10.55	0.97	32.76	1.95	40.81	3.09
Figure 4: RLA arrangement results with Small Magic-Plan (top) and Small RPLAN (bottom).
Figure 5: RLA arrangement results with Full MagicPlan dataset.
Our method outperforms all the others in all the settings and metrics, except the GED metric for Small MagicPlan, where the existing state-of-the-art Shabani et al. achieves a slightly better score. Their method enumerates all possible candidate arrangements by matching doors and evaluates the realism of each arrangement one by one. Their run-time is exponential in the number of rooms, and the system involves many heuristics. TransVector is end-to-end and achieves comparable performance with Shabani et al.; in fact, a much better MPE score for Small RPLAN. However, our system performs much better in every metric, demonstrating the power of Diffusion Models even for non-generative tasks, in our case, room layout arrangement. Figure 4 compares Shabani et al., Transvector and ours again on the small datasets. Figure 11 shows our results for the most challenging setting (i.e., Full MagicPlan dataset) comparing to Transformer with vector representation.
Table 2 presents a comparative analysis for CJP and VJP at three different noise levels. We have trained our model only for a noise-free case (i.e., 0 noise level) and used for the other two levels. Our approach consistently outperforms the current state-of-the-art. The performance gap is significant with the presence of noise, because Harel et al. relies on critical assumptions that matching edges have the same distance and neighboring angles sum to  $180^{\circ}$ . They have some tolerance to cope with noisy inputs but do not do well against our learning-based approach. Figs. 6, 7, and 8 show sample arrangement results for apictorial CJP, pictorial CJP, and VJP, respectively, supporting the above quantitative results. Note that the table shows the numbers for apictorial CJP. Due to the time complexity of Harel et al., we have only calculated the metrics on 20 samples for pictorial CJP, showing the same trend in results. The full details are referred to the supplementary, which also contains experiments on duplicate or missing pieces, where our method is surprisingly robust.
5.4 Ablation studies
We choose the room layout arrangement task for further ablation studies on our method. Table 7 shows the contributions of our two attention mechanisms (P-SA, G-SA) and the door matching loss. G-SA provides communications between every pair of corners in a house, and its removal has the most impact on the performance. Removing P-SA or the matching loss also leads to a significant drop in both MPE and GED metrics. Our method with all the components achieves the highest performance.
Table 2: CJP and VJP quantitative results with three metrics. Our method consistently shows superior performance with significant margins under the presence of noise. The second row shows ours without the matching loss. Harel et al. is unable to handle VJP due to the assumptions (See 5.2).
Dataset	Cross-cut Jigsaw Puzzle	Voronoi Jigsaw Puzzle
Metric	Overlap (↑)	Precision (↑)	Recall (↑)	Overlap (↑)	Precision (↑)	Recall (↑)
Noise Level	0	1	2	0	1	2	0	1	2	0	1	2	0	1	2
Harel et al.	0.91	0.70	0.30	0.95	0.77	0.33	0.99	0.78	0.30	×	×	×	×	×	×
Ours (w/o Lmatch)	0.91	0.90	0.89	0.94	0.92	0.90	0.77	0.76	0.73	0.65	0.64	0.63	0.75	0.71	0.53
Ours	0.94	0.94	0.93	0.97	0.95	0.93	0.91	0.92	0.91	0.70	0.68	0.67	0.78	0.75	0.74
Figure 6: Apictorial CJP arrangement results. Both methods work well in a noise-free case (left), but only our method maintains the performance with the presence of noise (right).
Table 3: Contributions of our two attention mechanisms (P-SA, G-SA) and the door matching loss  $(L_{\mathrm{match}})$ . Full MagicPlan is used.  $\checkmark$  indicates the feature being used. In case of  $L_{\mathrm{match}}$  "Doors" means matching loss has been applied only on door corners and "All corners" means matching loss has been applied to all corners including door corners.
P-SA	G-SA	Lmatch	MPE (↓)	GED (↓)
None	✓	None	48.2	4.9
✓	✓	None	43.4	4.5
✓	None	All corners	55.3	9.4
None	✓	All Corners	45.1	4.6
✓	✓	All corners	41.8	3.6
✓	None	Doors	56.9	9.3
None	✓	Doors	45.2	4.3
✓	✓	Doors	40.8	3.1
Table 4: Effects of the room-type (R-type) and the Door information. Full MagicPlan is used.  $\checkmark$  indicates the information being used. When a room-type is not used, we set a zero vector as a room-type one-hot vector, when room type is noisy we assign each room with a random room type. When the door information is not used, we do not pass the door-corner nodes to the network.
Train	Test	MPE(↓)	GED (↓)
R-Type	Door	R-Type	Door
✓	✓	None	✓	48.4	3.8
None	✓	None	✓	47.4	3.6
✓	✓	Noisy	✓	49.7	4.0
Noisy	✓	Noisy	✓	49.3	3.9
✓	✓	✓	None	46.9	5.2
✓	None	✓	None	45.5	5.2
✓	✓	✓	✓	40.8	3.1
Our layout arrangement uses a redundant representation (i.e., all the corners store the position and the orientation estimates of a piece), enriching the capacity and enabling direct communications between room/door corners (Sect. 4.1). To assess the effects of this redundancy, we have created a variant of our system with a compact representation, that is, each room/door has only a single node estimating a single copy of the room position and the orientation. We aggregate corner coordinates into a single embedding vector and pass as a condition (See supplementary for the details). The MPE/GED metrics for Full MagicPlan change from (41.23/3.16) to (51.62/5.52), a significant performance drop showing the importance of our redundant representation.
Figure 7: Pictorial CJP arrangement results.
Figure 8: VJP arrangement results.
A popular approach in the RLA literature is to align door detections/annotations to enumerate arrangement candidates Lambert et al. (2022a); Shabani et al. (2021). Our approach directly learns to infer the arrangement without relying on such "hard" heuristics. To further demonstrate the power of our approach, we remove or alter the room type and the door information during training and testing, where the current state-of-the-art methods Lambert et al. (2022a); Shabani et al. (2021) completely fail. Table 8 shows some performance drop, but the effects are marginal. Our numbers are still much better than TransVector with the full information (MPE=53.11 and GED=6.41 in Table 1), the only competing method capable of handling this most challenging setting. To evaluate the capacity of our network to handle overlaps, we conducted an experiment. In this experiment, we augmented our train/test dataset by randomly selecting up to two rooms per house. For each selected room, we employed one of the following strategies: 1) Duplicating the room with the ground truth (GT) room types. 2) Duplicating the room with random room types. 3) Enlarging the room by  $20\%$  to create partial overlaps. The dataset used for this experiment is Full MagicPlan dataset. Our evaluation metrics, MPE/GED, resulted in 48.3/3.9, respectively. Please also see the supplementary document and the video for more results, more visualizations, and animations of the denoising process.
This paper introduced an end-to-end neural architecture for spatial puzzle solving tasks. The proposed approach is faster, robust to data corruptions, end-to-end, and far superior to existing methods in all the metrics in the variety of tasks, namely Cross-cut jigsaw puzzle (pictorial and apictorial), Voronoi jigsaw puzzle, and room layout arrangement. Despite numerous advantages, our approach faces certain limitations. The primary challenge is its substantial need for training data, which limits our ability to process smaller datasets. For the room layout arrangement task, the utilization of raw image information could further enhance performance as discussed in Shabani et al. (2021); Lambert et al. (2022a). However, this would significantly increase the amount of data transfer from a mobile device, where on-device processing would become desirable. One key future work is the development of data-efficient (at training) and computation-efficient (at testing) neural architectures Kitaev et al. (2020). To our knowledge, this paper is the first to demonstrate that Diffusion Models, generally regarded as powerful generative models, are also effective in solving various challenging spatial arrangement tasks. The paper has the potential to motivate other researchers to further expand the applicability of emerging Diffusion Models, moving beyond content generation and into a myriad of other tasks.
Acknowledgment: This research is partially supported by NSERC Discovery Grants with Accelerator Supplements and the DND/NSERC Discovery Grant Supplement, NSERC Alliance Grants, and the John R. Evans Leaders Fund (JELF). We are also thankful to Magicplan for sharing the datasets.
Omri Avrahami, Ohad Fried, and Dani Lischinski. 2022. Blended Latent Diffusion. arXiv preprint arXiv:2206.02779 (2022).
A Methods details of our system and competing methods
We benefit from Transformers in our task in two ways. First, Transformers provide the capability of processing sequences with different lengths, which we use to process different number of room layouts/corners in the houses. Second, we utilize the self-attention module of Transformers to create optimal interaction and information-sharing among input tokens. These two features make Transformers an ideal backbone for our model. Our method uses six Transformer encoder blocks, and attention in each block has four heads. We also use an MLP For converting 256D Transformer output to rotation and position (4D). To keep the experiments fair, we use the same architecture for our transformer baselines as much as possible. In the following, we provide details corresponding to each of the baselines.
Transformer with a raster representation (TransRaster) uses the raster images to represent the input room layouts/types and the output room positions. Note that this baseline does not handle rotations as explained below. An input room layout is represented as a 20-channel  $256 \times 256$  semantic segmentation image, where there are 20 room/door types. The room center is aligned with the center of an image. An output room position is represented as a  $256 \times 256$  room occupancy image, which is ideally a translated version of the input room segmentation image at the correct room location. Given an output room occupancy image, we perform an exhaustive search over the possible room translations and find one with the most overlap between the occupancy image and the translated room segmentation image.<sup>3</sup> We use VisionTransformer Dosovitskiy et al. (2020) with a CNN decoder that takes a set of input room segmentation images and produces a set of room occupancy images.
In the other word TransRaster uses an encoder part of U-Net, which has 8 down-sampling blocks, converting each input room layout to a feature map of dimension 512. Each feature map (corresponding to a room layout) will become one input token for the Transformer. Input sequences length is equal to the number of rooms in a house, and information is shared among different rooms. We use six Transformer encoder blocks, and attention in each block has four heads. We pass the output of Transformer to a U-Net up-sampling model with eight Up-sampling blocks to change the dimension from 512 to  $256 \times 256$ .
Transformer with a vector representation uses the same backbone as our method; a linear layer converts the 28D input vector (i.e., 2 for the original corner coordinate and 20 for the room/door type one-hot vector) to the 256D feature map, six Transformer encoder blocks, and the attention in each block have four heads. We also use an MLP to convert 256D output embedding to 4D output.
Diffusion model, one room per node encodes each node as corresponding to a room instead of a corner in the room. To ease the implementation, we set the maximum number of nodes per room to 20 and we pad extra nodes when the room has less than 20 nodes with 0. We flatten the conditions per room and then use a linear layer to convert it to a 256D embedding vector. Each feature map represents a room and an input token for Transformer, we use the same Transformer as our method. After the Transformer blocks, a linear layer converts 256D output to 4D (i.e., 2 for the position and 2 for the rotation).
Shabani et al. Shabani et al. (2021) takes the input layout of each room with the resolution of  $256 \times 256$  with the same number of channels as the number of room types to pass each pixel as a one-hot vector of the corresponding room type. We use the same model as Shabani et al. (2021) and change the number of input channels to 11 for RPLAN and 20 for MagicPlan. To generate the arrangement candidates, we use the given room layouts of our dataset to connect doors, while we also use overlap filtering to reduce the number of candidates. Note that our datasets is significantly larger than the one in Shabani et al. (2021), enabling us to randomly select a positive or a negative candidate in each iteration and therefore remove the class imbalance weight used in Shabani et al. (2021). During the training for each house, we randomly select a GT with the label 1 or a faulty candidate with the label [0, 1) based on the number of mismatched doors. During the test, we pass all the possible candidates of each house and select the candidate with the highest score as the final prediction.
<sup>3</sup>We could expand the search space with possible room rotations, but rooms are often symmetric. To be simple, we use this baseline only for experiments when ground-truth rotations are given.
Harel et al. Harel and Ben-Shahar (2021) proposed a two-step algorithm for CJP. Their method considers two types of constraints to find plausible mates based on the length and angle of different pairs of connections. By estimating the matings hierarchically using these constraints, they approach the problem of finding positions as a multi-body spring-mass system. We utilize the authors' provided implementation<sup>4</sup> for comparison with our method. With a test dataset of 1000 crossing-cut puzzles, we restrict the running time of the spring-system algorithm to 2 minutes per puzzle. Furthermore, unlike the provided implementation, we also consider failure cases in the metrics. Regarding the pictorial case, the authors score a candidate mating by extrapolating the images of puzzle pieces and considering the difference of the mean color value on the edges. We do not impose any time limit as we evaluate only on 20 samples.
B Datasets Details and Preprocessing
Room Type	3	4	5	6	7	8	9	10	All
Master bedroom	0.20	0.26	0.29	0.32	0.39	0.49	0.49	0.53	0.34
Living room	0.52	0.56	0.59	0.65	0.71	0.75	0.79	0.79	0.65
Kitchen	0.42	0.47	0.52	0.59	0.68	0.71	0.76	0.79	0.59
Bathroom	0.54	0.70	0.85	0.96	1.12	1.28	1.33	1.47	0.96
Toilet	0.07	0.11	0.15	0.22	0.22	0.25	0.27	0.26	0.18
Corridor	0.07	0.12	0.15	0.19	0.25	0.32	0.37	0.45	0.21
Closet	0.13	0.18	0.22	0.32	0.48	0.68	0.92	1.22	0.41
Hall	0.35	0.55	0.68	0.77	0.85	0.91	0.98	1.08	0.73
Laundry room	0.05	0.05	0.05	0.07	0.10	0.13	0.18	0.23	0.09
Bedroom	0.34	0.69	1.08	1.32	1.51	1.74	1.93	2.10	1.23
Balcony	0.05	0.08	0.16	0.32	0.40	0.48	0.56	0.64	0.29
Dining room	0.13	0.12	0.12	0.14	0.17	0.19	0.22	0.25	0.15
Private office	0.00	0.01	0.02	0.05	0.05	0.07	0.08	0.10	0.05
Den	0.05	0.05	0.06	0.7	0.09	0.11	0.12	0.12	0.08
Storage	0.00	0.01	0.01	0.02	0.02	0.02	0.03	0.03	0.02
Others	0.00	0.01	0.01	0.02	0.03	0.04	0.04	0.07	0.03
Doors	2.84	3.82	4.82	5.43	6.92	8.02	6.10	10.18	5.90
Table 5: MagicPlan dataset consists of floorplans with 3 to 10 rooms. The table shows average number of rooms with a specific room type based on the total number of rooms in the house.
We normalize the puzzles/floorplans for each task and dataset by scaling them to fit within a  $1 \times 1$  square, and we also resize all corresponding images to dimensions of  $256 \times 256$  in the case of pictorial CJP. While this normalization process does not introduce any essential additional information during testing in CJP and VJP, it could potentially enable the network to cheat in RLA, as the longer extent of arranged floorplans is always fixed to 1. To address this issue, during testing in RLA, we apply a random scaling factor in the range of [0.8, 1.0] to the room shapes of each house. In the subsequent sections, we provide a detailed description of each dataset. In the following, we provide additional statistics for our floorplan datasets.
The Voronoi Jigsaw Puzzle dataset consists of 200k training puzzles and 1k testing puzzles. These puzzles were created by randomly selecting 3 to 15 points, The individual pieces of the puzzles were obtained by extracting the Voronoi cells corresponding to these points. There are 1,066, 14,033, 23,715, 20,279, 16,073, 15,235, 16,428, 15,018, 16,318, 15,096, 16,487, 16,233, and 15,670 puzzles with 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, and 15 pieces respectively. Each piece has a minimum, maximum, and average of 3, 20, and 4.51 corners respectively. The minimum, maximum, and average number of corners per puzzle are 10, 93, and 42.24.
Cross-cut Jigsaw Puzzle (CJP) are consist of 100k training and 1k testing puzzles, where each one were generated using Harel and Ben-Shahar (2021) method which generate a convex polygon and cuts it by 3 to 5 lines. There are 1719, 6046, 15854, 14521, 6905, 10929, 12065, 8361, 6521, 8192, 7663, 4642, 1508, 73, and 1 puzzles with 3, to 18 pieces respectively. Each piece has a minimum,
4https://icvl.cs.bgu.ac.il/polygonal-puzzle-solving/
maximum, and average of 3, 13, and 4.47 corners respectively. The minimum, maximum, and average number of corners per puzzle are 16, 76, and 41.99.
MagicPlan dataset consists of roughly 98K houses/apartments, which we divide into 93K training and 5K testing samples. The number of rooms in a house ranges from 3 to 10. Concretely, 11661, 16322, 19171, 17582, 13200, 9649, 6780, and 4415 houses contain 3, 4, 5, 6, 7, 8, 9, and 10 rooms, respectively. The minimum and maximum numbers of corners in a house are 12 and 182. Table 5 shows average number of rooms with a specific room type based on the total number of rooms in the house.
In the RPLAN dataset, we divide 60K samples in RPLAN to 55K train and 5K test. The number of rooms in a house ranges from 3 to 8. Concretely 99, 582, 5083, 19551, 21921, and 13235 houses contain 3, 4, 5, 6, 7, and 8 rooms, respectively.
C Additional ablation studies
C.1 Additional ablation studies on room layout arrangement
Table 6: Main quantitative results with two metrics: Positional Error (MPE) and Graph Editing Distance (GED). This table show a case where the ground-truth rotations are given, as TransRaster baseline cannot handle rotations. Small RPLAN (resp. Small MagicPlan) is a subset of the corresponding full dataset, consisting of houses with at most 6 rooms. The small datasets are created for Shabani et al., which is not scalable to many rooms. Our method is stochastic and shows both the mean and the standard deviation.
Dataset	Small RPLAN	Full RPLAN	Small MagicPlan	Full JigsawPlan
Metric	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)	MPE (↓)	GED (↓)
Shabani et al.	17.6	1.0	X	X	32.2	1.1	X	X
TransRaster	13.9	1.2	15.7	2.1	36.1	2.1	41.9	4.1
TransVector	12.9	1.1	13.9	2.0	37.7	1.9	42.8	4.0
Ours	4.6±0.7	0.4±0.0	5.4±0.7	0.6±0.0	17.5±0.8	1.0±0.4	27.9±0.7	2.7±0.5
Figure 13 shows the raw estimated position information at each room/door corner before the room-wise averaging. Since the ground-truth has the same pose parameters for all corners in a room/door, the network learns to produce consistent parameters. Figure 14 shows five pose estimation results by our system while varying the initial noise  $x_{T}$ . While there are minor differences, the overall room arrangements are similar and close to the ground-truth, indicating that the Diffusion model is capable of producing consistent results given enough constraints as a pose estimation system, as opposed to a generative model whose original goal is to create a diverse set of answers.
Figure 9: Visualization of predicted layouts at step "t"s. At  $t = 1000$ , position parameters at each corner are initialized by a Gaussian noise, and at  $t = 0$ , there is the final predicted layout. The top row shows the predicted layout without averaging/voting, and the bottom row shows with averaging/voting. To make it more clear, we show doors by their corners.
Figure 10: Qualitative evaluations of our approach against the three competing methods. The top two rows are from Small MagicPlan. The bottom two rows is from Small RPLAN. The GT rotations are given for all the cases to enable comparisons with all the methods.
C.1.1 Additional ablation on RPLAN
The main paper shows the ablation studies on the MagicPlan dataset in case of room layout arrangement task. This part of supplementary will present the same study results on RPLAN dataset. Table 7 shows the impact of our attention module and door matching loss on performance and Table 8 shows the impact of noise in the room type and door detection on our performance, although there is a performance drop, our method still works better than the competing methods.