Puzzlefusion: Unleashing the Power of Diffusion Models for Spatial Puzzle Solving
Authors: Sepidehsadat Hosseini
Mohammad Amin Shabani
Saghar Irandoust
Yasutaka Furukawa
Venue: NeurIPS 2023 spotlight (2023)
Keywords: Diffusion
Jigsaw
puzzle solving
TL;DR: This paper presents an end-to-end neural architecture based on Diffusion Models for spatial puzzle solving, particularly jigsaw puzzle and room arrangement tasks.

Abstract: This paper presents an end-to-end neural architecture based on Diffusion Models for spatial puzzle solving, particularly jigsaw puzzle and room arrangement tasks.
In the latter task, for instance, the proposed system ``PuzzleFusion'' takes a set of room layouts as polygonal curves in the top-down view and aligns the room layout pieces by estimating their 2D translations and rotations, akin to solving the jigsaw puzzle of room layouts. A surprising discovery of the paper is that the simple use of a Diffusion Model effectively solves these challenging spatial puzzle tasks as a conditional generation process. 
To enable learning of an end-to-end neural system, the paper introduces new datasets with ground-truth arrangements: 1) 2D Voronoi Jigsaw Dataset, a synthetic one where pieces are generated by voronoi diagram of 2D pointset; and 2) MagicPlan Dataset, a real one from a production pipeline by MagicPlan, where pieces are room layouts constructed by augmented reality App by real-estate consumers.
The qualitative and quantitative evaluations demonstrate that the proposed approach outperforms the competing methods by significant margins in all three spatial puzzle tasks. We have provided code and data in https://sepidsh.github.io/puzzlefusion.

PuzzleFusion: Unleashing the Power of
Diffusion Models for Spatial Puzzle Solving
Sepidehsadat Hosseini, Mohammad Amin Shabani, Saghar Irandoust", Yasutaka Furukawa
Simon Fraser University (sepidh, mshabani, sirandou, furukawa @sfu..a
Abstract
This paper presents an end-to-end neural architecture based on Diffusion Models for spatial puzzle solving, particularly jigsaw puzzle and room arrangement tasks. In the latter task, for instance, the proposed system takes set of room layouts polygonal curves in the top-down view and aligns the room layout pieces by estimating their 2D translations and rotations, akin to solving the jigsaw puzzle of room layouts. A surprising discovery of the paper is that the simple use of a Diffusion Model effectively solves these challenging spatial puzzle tasks as a conditional generation process. To enable learning of an end-to end-to-end neural system, the paper introduces new datasets with ground-truth arrangements: 1) 2D Voronoi jigsaw dataset, synthetic one where pieces are generated by Voronoi diagram of 2D pointset; and 2) MagicPlan dataset, a real one offered by MagicPlan from its production pipeline, where pieces room layouts constructed by augmented reality App by real-estate consumers. The qualitative and quantitative evaluations demonstrate that our approach outperforms the competing methods by significant margins in all the tasks. We have provided code and data here.
Introduction
Spatial puzzle solving demands meticulous reasoning and arrangement of elements within given space. The classic example of jigsaw puzzles, which many of us have enjoyed as a recreational ac- tivity, showcases the challenge and satisfaction of such tasks. Applications of jigsaw puzzles extend beyond mere entertainment, encompassing areas such as the restoration of shattered 2D artwork and documents Das all (2017)); McBride Kimal (НEE), image stitching in computer vision Hame mondeh and Polleti (201))) and even DNA sequence assembly in genomics Pop (2(19); Marand and Burged (20077) In real estate, room layout arrangement has emerged compelling spatial puz- zle task, where consumers leverage mobile devices to scan individual rooms that are to be arranged into a floorplan.
Spatial puzzle solving poses a considerable challenge even for humans, making it an engaging men- tal exercise. Addressing this challenge for computational approaches is even demanding, de- spite the emergence of deep learning. Current state-of-thh-aan techniques typically enumerate pairs of aligned pieces, evaluate their compatibility by learning, and employ optimization search meth- ods to identify the most likely global arrangement Harel and Ben-Shanar (2021); Hoff and Olver (2014a); Shih and LiI (20I). However, these approaches struggle to scale the complexity of global arrangement increases exponentially with the number of pieces.
This paper makes breakthrough in spatial puzzle solving with end-to-end neural architecture based on Diffusion Models Ho et al (2020), named PuzzleFusion. The surprising discovery this paper is that Diffusion Model, typically regarded as a powerful generative model, is an effective
Work done while being student SFU
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Ground Truth
Living Room
Balcomy
Hall
Room Layouls
Bedroom
Bathroom
Door
Kitchen
Closet
Initialization
Final Prediction
Figure I: Room layout arrangement the task of taking set of room layouts and their correspond- ing room types the input and predicting the position and the orientation of each room. The biggest discovery and surprise of this paper is that conditional generation by a Diffusion Model solves this challenging problem.
spatial puzzle solver. PuzzleFusion formulates the spatial puzzle as conditional generation process where the piece information the condition.
Concretely, each puzzle piece is polygonal shape, represented by sequence of 2D corner coor- dinates. A piece associated with task-dependent properties, such as texture image for pictorial jigsaw puzzles or the room types and door locations for Room Layout Arrangement. The output 2D position and an orientation for each piece. The forward diffusion process adds noise feature vector initialized with the ground truth position and orientation. The reverse denoising pro- cess learns infer position and orientation subject to the piece shape and property information as condition.
Qualitative and quantitative evaluations over the three tasks (Cross-cut jigsaw, Voronoi jigsaw, and room layout arrangement) demonstrate that PuzzleFusion outperforms the current state-of-the-art methods by significant margins. The paper also makes dataset contributions by introducing large- scale datasets with ground-truth labels for Voronoi jigsaw and room layout arrangement tasks. Specifically, the room layout arrangement dataset comes from production pipeline by MagicPlan (https:/www.maagg www.:w/w.ppppa...ommom.om appp/ consisting of room-layouts and floorplans for 98,780 houses, which we obtained permission to share with the research community.
In summary, this paper makes the following three key contributions: 1) An end-to-end neural archi- tecture based on Diffusion Models for spatial puzzle solving; 2) State-of-the-ar performance across three spatial puzzle tasks in terms of accuracy and speed; and 3) The new spatial puzzle datasets including room-layouts and floorplans for 98. 780 houses from a production pipeline. We will make all our code and data public.
Related Work
Spatial puzzle solving is closely related to Structure from Motion (SfM), pose estimation, arrange- ment learning, and more. The section discusses the related techniques.
Feature matching has been successful for the SfM problem Snavely et all (2006); Liet all (2020);
robust feature matching by leam- Lin ell (2016). The rise of deep neural networks enables ing Yi all (2018); Sun et all (2022);; 'aarlnn et all (202(). However, these techniques require visual overlaps. Our task has little to visual overlaps between adjacent images.
Geometry inference estimates relative pose between images or partial scans with minimal over- laps by registering inferred or hallucinated geometry. A popular approach learns the priors of room shapes and alternates pairwise alignment and scene completion et (2019); Yang all (2019).
Yang et al. Yang al. (2020) combines global relative pose estimation and local pose refinement

---

with panoramas. These techniques learn priors of a single room, while our approach learns the arrangements of multiple rooms on house scale.
Arrangement learning is the current state-of-the for indoor room layout arrangement. An early work uses windows align indoor and outdoor reconstructions Cohen et all (20T6). Shabani et al. Shabani et all (2021T) use doors to enumerate room arrangements and learn to score each can- didate. Their approach is exponential in the number of rooms with many heuristics. Lambert al. Lambert et all (200222)) uses doors, windows, and openings to create room alignment hypotheses. They utilize depth maps create top-down views and learn verify the correctness, improving a run-time from exponential to polynomial. Our end-to-end approach does not enumerate arrange- ment candidates and makes significant performance improvements. Lastly, an annotated site map and SIM reconstructions are aligned to solve challenging structure from motion problem Mariis
Brualla (2014);; Hosseini and Furukaws (2022), which they coined as a "3D jigsaw puzzle".
Puzzle Solving has been an engaging research area for long time Ereeman and Garden (M))
Radack and Badler (1982); Markakı and Panagiotakıs (2022), ranging from pictorial puzzles with image information and (2013) Le and (2019); Foler-Franklin et all (2010) to apictorial puzzles with only geometry information Goldberg et all (2002); Harel and Sen-Shabar (2020); Hoff and Olver (20)ER, 20T3); Harel and Ben-Shabar (2022 l). In previous studies, heuristic based meth- ods, such edge and color matching Wolfson all (1998) Nielsen (2008); Chung et all (T998), have been predominantly used for solving both methods. Recently, deep learning-basee seek learn high-level pictorial features Noroozi and Favard (2016); | et all (2021a). However, they suffer from poor geometric reasoning and are limited to simple square puzzles. Consequently, researchers combine learning-basse models with heuristics handle more complex puzzles Lecand Li (2) Nonetheless, these methods still require explicit pairwise comparison of pieces and lack data-driven high-level reasoning Our method overcomes these limitations through innovative end-to-end use of Diffusion Models which are usually regarded as powerful generative models.
methods
Diffusion models (DMs) emerging generative models, which slowly corrupt sample by adding noise Ho et all (2020); Dhariwal and Nichol (2021)); Nichol and Dharwall (202TE); Shabanı et all (2022), learn invert the process, and generate diverse set of samples from noise signals. DMs
have established SOTA performances in numerous tasks such as image colorization/mpainiing Song at (2020;; Nichol ala (202T), image to image translation Sasaki all (2021); Zhao all (2022).
text to image Ramesh et all (2(222), super-resolution Rombach el all (2021)) Saharia et all (2(22));
Li al (20216), image and semantic editing (Arrhami all (2022) ; Meng all (2022). and de-
noising Kawar et all (2022)). Recent works use DMs as representation learners for discriminative tasks such image segmentation Baranchuk et al (20Z4T); Wolleh et (202T)). Diffusion inspired models have been used for human pose estimation Shan ctall (2023); Gong ef all (20023) and object
placement Wei all (2023) however, both tasks do require sophisticated shape comparisons. While in our case more advanced and intricate approach is necessary to capture and exploit all details of the shapes.
Crosecait (AP)
Crosecnt (P)
MagicPlan
Master Bedroom
Kitchen
Living Room
Study
Dining
Laundry
Sterage
Bedroosa
Closet
Ceeridor
Hall
Balcony
Figure 2: Spatial puzzle datasets for Voronoi and Cross-cut puzzle solving (Left) and room layout arrangement (Right). We consider both pictorial and apictorial versions of the Cross-cut jigsaw puzzle. Note that colors in puzzle solving problems are random and just indicate different pieces.
3

---

3 Spatial Puzzle Solving Tasks
Spatial puzzle solving involves arranging puzzle pieces, each with its geometry and optional features such as images or categories. In real scenarios, puzzle pieces are affected by erosion, duplication, or loss. This paper investigates three distinct puzzle-solving tasks.
Cross-cut Jigsaw Puzzle (CJP) generates pieces by making random straight cuts through larger polygonal shape Harel and Shaaat (2021).. Puzzle pieces are convex polygons, each with an arbitrary number of neighboring pieces. The sum of the two adjacent angles at any corner equals 180". The application of the technique could be the restoration of shattered artwork.
Voronoi Jigsaw Puzzle (VJP) generates pieces by randomly sampling points within predeter- mined bounding box and extracting the cells of the Voronoi diagram as the pieces. With the lack of the "180" constraint", VJP significantly more challenging, where no effective solution has been presented in the literature to our knowledge. Voronoi diagrams play role in biological systems like cell arrangements Bock et all (20IC), potentially useful studying natural phenomena.
Room Layout Arrangement (RLA) determines the room arrangement and the corresponding floor- plan, offering key application in real estate. Contrary to the prior tasks, a corner may align along the edge of another piece, expanding the solution space to be explored. Pieces come from the room layout estimation algorithms Shabami all (202II); Lambert et all (2022a) or interactive augmented reality apps used by consumers.
Task input/output:
The input is set of N polygons (puzzle pieces in CJP/VJP and room layouts in RLA), each of which is a sequence of corner coordinates forming a ID polygonal loop. In RLA, door piece is also given line segment with two corners, and room piece is associated with room type as a 20D one-hot vector. For a pictorial version of CJP, a piece is associated with a 128D image feature vector obtained from pretrained auto-encoder. For simplicity, and door-comers, and use C to denote the 2D coordinate of the ith corner in the rth polygon. T denotes the image feature the room type vector, Please refer supplementary for details. The output is the position of the piece/room center and the rotation around (an angle between 0 and 2m for CJP/VJP and 4-fold Manhattan rotation for RLA) The center is the average of the corners.
mix room-corners
For CJP and VJP, we adopt metrics used in previous work Harel and Ben-Shahail (2021), namely Overlap, Precision, and Recall. The overlap score is the average IoU of pieces with the ground truth. Precision and Recall are on the connectivity of neighboring pieces. For RLA, we con- sider two metrics. The first metric is the Mean Positional Error in pixels (MPE) over the rooms Shaa bani el all (2(12I). The second metric evaluates the correctness of the room connectivity in the reconstruction. We borrow Graph Edit Distance (GED) by Nauata et al. INuuata et all (202I), which counts the number of user edits necessary to fix the connectivity graph. We declare that two rooms connected if the door centers are within 5 pixels from the two rooms. Note that have designed task-specific metrics, respecting the methodologies of existing literature while providing a thorough evaluation of our performance.
Metries:
Spatial Puzzle Solving as Conditional Generation
Our idea is simple, using Diffusion Model to 'conditionally generate" the correct arrangement, where the input the center and rotation of room layouts/puzzle pieces and their types and shapes are the conditions. This section explains the forward and the reverse processes.
4.1 Forward process
The forward process adds Gaussian noise to an arrangement. A compact representation would be per polygon positions and rotations. Instead, we will use a redundant representation, where a center position and rotation are stored at each corner.
"Shabani et al. Shaban all (002)) used a permissive metric (the availability of the "correct" solution in the k results with a certain cror-tolerance as the task was challenging. We make great improvements and use standard metric.

---

( room index
corner index
Lmaach
Encoder Layers
Feature iMhhedding of Comers
Poly.
MLP of Corners
Global SA
Add & Norm
It
H--1
Poly:
Door
Corner
Forward
Reverse
TT
Figure 3: Illustration of our diffusion model architecture employed for the RLA task. Given the arrangement estimation I; = {xT,,), the reverse process infers the noise {or,) and recovers It-1 { while injecting the original room shapes {CT} and types {7"} as the condition. Each room corner holds the room position and rotation estimation. The reverse process starts from øT and denoises towards zo.
There are a few reasons. Our reverse process is based on a Transformer architecture where each po- sition/rotation estimation becomes node. Our approach 1) enriches the capacity of the urrangement representation (also an adaptive capacity, where a complex polygon with more corners is given more capacity); 2) allows direct communications between corners/doors for which will have specific loss; and 3) makes it straightforward to combine with the condition (i.e., original corner coordinates and room types).
Concretely, we use "i,, to denote the position/rotation of the rth room/piece stored at the ith corner at time & of the diffusion process, where í varies from 0 to 1,000 in our experiments:
t (pL of,))
pf, and of,, denote the polygon-center position (a 2D vector) and the rotation. We consider rotation as 2D vector obtained from rotation matrix including cos(o,,) and -sin(of,). The forward process adds a noise by sampling &. (O, I) with a standard cosine noise scheduling with variance ax) Nichol and Dhariwal (20216):
f(0)* f(t) = cos 1/T 0.008 π
f(1)
x = Vartl, + vI- a,d ā
(2)
1.0.008 0.008 2
4.2 Reverse process
Figure E iillustrates our reverse process for room layout arrangement, which takes the arrangement
{ time and infers the noise of under the condition of the original room shapes {CT} (i.e..
coreer position with respect the room center) and the room types as one-hot vector {
Feature embedding: The reverse process based on a Transformer architecture where every corner node. We initialize its feature embedding
N Lin(x{,) + Lin((CT, r, i]) + MLP(!) + Lin(7").
(3)
The first term uses linear layer to convert 4D vector (i.e., 2 for the room center coordinate and 2 for the rotation vector) to 256D embedding vector. The second term also uses linear layer to convert condition vector (i... 2 for the corner coordinate respecting to the center of the polygon

---

in addition to 32 for polygon index which shows which polygon each comner belongs r and 32 to show the corner index in each polygon i), The third term uses 2-layer MLP to convert a time step I to 256D vector. The fourth term is extra information such as room/door type (20 types) for RLA which we extract it by applying a linear layer on 20D one-hot vector or 128D for image features in pictorial CJP)
Attention modules: Feature embeddings {&C,) } go through six blocks of self-attention modules that have two different attention mechanisms: Polygon Self Attention (P-SA) and Global Self Attention (G-SA). P-SA limits pairwise interactions between corners in each polygon. P-SA is akin sparse self attention family Clll (2011); Child (201); al (2011) which helps to generate consistent positions and rotations at different corners of polygon. G-SA is standard self-attention
between all corners of a puzzle or a house. After the attention blocks, linear layer converts 256D
embedding back to 4D representation o7 which is used for the following denoising formula Ho
(2020):
1-a
1-
2
VI-ai a
+ V1 0'2..
Van
VI
(N0 1) for & > 1 and otherwise O. For the final result at time t = 0, we take the average polygon center position and the polygon rotation.
Loss functions: There are two loss functions. We first follow all (2020); Dhariwal and Nichol (2021) and use a standard noise regression loss on 8:
Isimll or 110% 87,
(5)
To enhance the quality of supervision, we propose "matching" loss, specifically aimed at the vertices where incident edges meet. These vertices are shared by two polygons. We denote the indices of these two polygons as ri and r2, then the corresponding vertex indices within rı and r2 as i and is, respectively. The corresponding vertices must be at the same position, and as such, we calculate the loss as the Fuclidean distance between their coordinates:
Lmaatch = Ei
||
-
Cra 1
(6)
Cg ph + Roy CT,
oo ,,,))
= Jh ( VV
- ar0%1 /oo..
Rap, denotes the rotation matrix corresponding 0a For the room layout arrangement task, add the loss only doors without walls, which yields superior results in experiments. The total loss is defined as the sum of the above loss functions, Taonal = Tsimple + Tmmatch- In practice, we found that adding T-atthh only for & < 500 results in better performance.
5 Experimental Results and Discussions
We have implemented the system with PyTorch Paszke ef all (2019), using a workstation with 3.70GHz Intel 19-10X00X CPU (20 cores) and two NVIDIA RTX A6000 GPUs. We use the AdamW Loshchulov and Hutter (2017); Kingma and Ba (201) optimizer with 31 0.9,
82 0.999, weight decay equal 0.05, and the batch size of 512. The learning rate is initial- ized to 0.0005. It takes roughly 24 hours to train model and 3 seconds to estimate the arrangement
for one sample.
5.1 Preprocessing and Datasets
We have carefully pre-processed and prepared datasets for fair evaluation, whose details are referred supplementary. We here provide summary points for the three tasks (See Fig. D)
Cross-cut Jigsaw Puzzle (CJP): We have used code provided by Harel et al. Harel and Shahar (2021) to generate 100k/lk 1k puzzles for training/testing. The code generates convex polygon and

---

cuts it by 3 to 5 lines. For the pictorial version, we have used images from COCO 2017 dataset Lmn all (2013), that is, randomly selecting training/testing image of the dataset for each train- ing/testing CJP sample. Following Harel and Ben-Shahar (2021), to simulate real-world scenario, we perturb corner coordinates with noise with three different levels to create datasets under three different noise levels with thresholds equal to 0, 2, and 4.
We have generated 200k/1k puzzles for training/testing by randomly sampling 3 to 15 points inside a random rectangle and extracting their Voronoi cells as the pieces. We have created datasets with three different levels of noise. For each corner, we added a random number from a Gaussian distribution with o² equal to O, 2, and 4 for noise levels 0, I, and 2, respectively.
Voronoi Jigsaw Puzzle (VJP):
Room Layout Arrangement (RLA): MagicPlan (https://www. (utps://www.magiiqpan...p magicplan.app), a mobile software company for real estate and construction, agrees to share production data with us and the re- search community, where the paper introduces the MagicPlan dataset, containing room shapes and their ground-truth arrangement for 98,780 single-story houses/apartments. We split the data into 93,780/55000 5,00 training/testing samples. MagicPlan software reconstructs room shapes one by one by asking users to click room corners through an augmented reality app. Room shapes are Manhattan- rectified as an enforcement of the app, then manually arranged to form a floorplan, which we seek automate. Each room is associated with room type. The number of rooms (resp. corners) in house ranges from 3 to 10 (resp. 12 to 182). We also use RPLAN et (2019) for evaluation, containing 60k floorplans. We divide into 55k/5k training/testinn. where the number of rooms (resp. corners) in a house ranges from to 8 (resp. 14 to 98). Note that RICOH dataset Saharia et all (M024) and ZIND dataset Lamhert (2022) are too small for network training and are not used in our experiments.
5.2 Competing methods
Harel et al. Harel and Ben-Shahar (2021) and Shabani et al. Shabani et al. (2022) are stateoof-the-an 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 or corners may not meet. Shabani el al. 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-Shahai (2(21) proposed 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 spring system.
Shabani et al. Shabani et all (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.
Vector is a baseline Transformer network with vector representation that directly estimates the pose parameters instead iterative denoising. Trans Vector shares the same architecture as our denoising network at the core with the following changes: Remove time-dependent features (x²,, and t) from the embedding (L); 2) Change the supervision di 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 W compares the proposed approach against Shabani er al. Shabani et all (2021) and Trans Vector. Shabani el al. runs exponentially in the number of pieces (i.e., rooms), taking hours or even days to process a single house with seven more rooms. Therefore, we collect small houses (i.e., at most six rooms) to create "Small RPLAN" and "Small MagicPlan" datasets for evaluations. For each experimental setting (e.g., Small MagicPlan), we train network for each method.
7

---

Table 1: RLA quantitative results with two metrics: Positional Error (MPE) and Graph Editing Dis- tance (GED). Small RPLAN (resp. Small MagicPlan) a subset of the corresponding full dataset, consisting of houses with at most 6 rooms. The small datasets created for Shabani et al. Shhabani at (2021), which not scalable to many rooms.
Dataset
Small RPLAN
Full RPLAN
Small MagicPlan
Full MagicPlan
Metric
MPF (4)
GED (4)
MPE (1)
GED (1)
MPF (4)
GED (4)
MPF (L)
GRD (1)
Shabani et al.
1.28
36.63
1.89
Trans Vector
36.09 8.65
1.51
46.18
2.27
40.80
2.38
53.11
Ours
0.90
10.55
0.97
32.76
1.95
40.81
3.09
Inpux
Shabani al.
Trans Vector
Ouns
Ground Truth
Inpat
Trans. Vector
Ouss
Ground Truth
Figure 4: RLA arrangement results with Small Magic-
Figure 5: RLA arrangement results with Full MagicPlan dataset.
Plan (top) and Small RPLAN (bottom).
Our method outperforms all the others all the settings and metrics, except the GED metric for Small MagicPlan, where the existing state-of-the-e art Shabani er al. achieves slightly better score. Their method enumerates all possible candidate arrangements by matching doors and evaluates the by one. Their run-time is exponential in the number of rooms,
realism of each arrangement and the system involves many heuristics. Trans Vector is end-to-end and achieves comparable per- formance with Shabani er al.; in fact, much better MPE score for Small RPLAN. However, 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 E compares Shabani al., Transvector and ours again on the small datasets. Figure [ shows our results for the most challeng- ing setting (i.e., Full MagicPlan dataset) comparing Transformer with vector representation.
Table presents comparative analysis for CJP and VJP at three different noise levels. We have trained our model only for noise-free case (i.e., 0 noise level) and used for the other two levels. Our approach consistently outperforms the current state-of-the-arr, The performance gap is significant with the presence of noise, because Harel er al. relies eritical assumptions that matching edges have the same distance and neighboring angles sum to 180". They have some tolerance to cope with noisy inputs but do not do well against our learning-bass based approach. Figs. G. and E 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 el 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 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 most impact on the performance. Removing P-SA or the matching loss also leads to significant drop in both MPE and GED metries. Our method with all the components achieves the highest performance.

---

Table 2: CJP and VJP quantitative results with three metrics. Our method consistently shows supe- noor performance with significant margins under the presence noise. The second row shows ours without the matching loss. Harel al. is unable to handle VJP due to the assumptions (See 5.20.
Dataset
Cross-cut Jigsaw Puzzle
Voronoi Jigsaw Puzzle
Metric
Overlap (t)
Precision (t)
Recall (t)
Overlap (t)
Precision (t)
Recall (1)
Noise Level
0
2
0
2
0
0
0
2
0
2
Harel al.
0,91 0.70 0.30 0.95 0.77 0.33
0.99 0.78
0.30
X
X
X
Ours (wlo m
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.71
0.53
0.52
0.52
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
0.60
0.57
0.55
Noise Level 0
Noisc Level
Harel al.
Ours
GT
Harel et al.
Ours
GT
Figure 6: Apictorial CJP arrangement results. Both methods work well in noise-free case (left), but only our method maintains the performance with the presence of noise (right).
Table 4: Effects of the room-type (R-type) and the Door information. Full MagicPlan is used. indicates the information being used. When room-type not used, set zero vector a room-type one-hot vector, when room type is noisy assign each room with random room type When the door information is not used, we do not pass the door-comer nodes to the network.
Table 3: Contributions of
attention mechanisms (P-SA, G-SA) and the door match- ing loss mmmtch)) Full MagicPlan is used. Viddicates the feature being used. In case of matchh "Doors" means matching loss has been applied only on door corners and "All cor- ners" means matching loss has been applied all corners including door corners.
P-SA G-SA
Immatch
MPE (4)
GED (L)
Train
Test
MPF(J)
GED
None
None
48.2
4.9
R-Type
Door
R-Type
Door
V
None
43.4
4.5
None
48.4
3.8
None
All corners
55.3
9.4
None
None
47.4
3.6
None
All Corners
45.1
4.6
V Noisy
Noisy
49.7
4.0
All corners
41.8
3.6
Noisy
49.3
3.9
None
Doors
56.9
9.3
V
None
46.9
5.2
None
Doors
45.2
4.3
None
None
45.5
5.2
Doors
40.8
3.1
40.8
3.1
Our layout arrangement uses redundant representation (i.c., all the corners store the position and the orientation estimates of piece), enriching the capacity and enabling direct communications between room/door corners (Sect. I.. To assess the effects of this redundancy, we have created 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 orientation. We aggregate corner coor- dinates into single embedding vector and pass as condition (See supplementary for the details). The MPE/GED metrics for Full MagicPlan change from (41.23/3.1) 361 to (61.62/55.22 significant performance drop showing the importance of our redundant representation.

---

Harel et
Ours
Noūse Level 0
Noise Level -2
GFT
Ours
GT
Figure 7: Pictorial CJP arrangement results.
Figure 8: VJP arrangement results.
popular approach in the RLA literature is align door detections/amnotations enumerate ar- rangement candidates Lambert ct al (20227); Shabani (2022). 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 wtat-offthee-ar art methods Lambert all (2022a); Shabani all (2021) completely fail. Table N shows some performance drop, but the effects are marginal. Our numbers are still much better than Vector with the full information (MPE=53.1 1 and GED=6.4 41 Table II), the only competing method capable of handling this most challenging setting. To evaluate the capacity of our network handle overlaps, we conducted an experiment. In this experiment, we augmented train/test dataset by randomly selecting up to two rooms per house. For each selected room, 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% 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.
5.5 Conclusion
This paper introduced an end-to-end neural architecture for spatial puzzle solving tasks. The pro- posed approach is faster, robust data corruptions, end-to-end, and far superior 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, utilization raw image information could further enhance performance discussed in Shaban el all (2021)) ID); Lamberi aLal (20223). 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 develop- ment of data-efficient (at training) and computation-efficieet (at testing) neural architectures Kilacy atl (2020). To knowledge, this paper is the first to demonstrate that Diffusion Models, gener- ally regarded as powerful generative models, are also effective in solving various challenging spatial arrangement tasks. The paper has the potential to motivate other researchers further expand the applicability of emerging Diffusion Models, moving beyond content generation and into myriad of other tasks.
Acknowledgmentt:
This research is partially supported by NSERC Discovery Grants with Accelera- tor Supplements and the DND/NSERC Discovery Grant Supplement, NSERC Alliance Grants, and John Evans Leaders Fund (JELF). We are also thankful to Magicplan for sharing the datasets.
References
Omri Avrahami, Ohad Fried, and Dani Lischinski. 2022. Blended Latent Diffusion. arXiv preprint 02777 (2022).
10

---

Dmitry Baranchuk, Ivan Rubachev, Andrey Voynov, Valentin Khrulkov, and Artem Babenko. 2021. Label-Efficient Semantic Segmentation with Diffusion Models. (2021). P.Xi::e.CV//212..000
Martin Bock, Amit Kumar Tyagi, Jan-Ulrich Kreft, and Wolfgang Alt. 2010. Generalized voronoi tessellation as a model of two-dimensional cell tissue dynamics. Bulletin of mathematical biology 72 (2010), 1696-1731
Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. 2019. Generating long sequences with sparse transformers. arXiv preprint arXiv: 10509 (2019).
Min Gyo Chung, Margaret M Fleck, and David A Forsyth. 1998. Jigsaw puzzle solver using shape and color. ICSP'98. 1998 Fourth International Conference on Signal Processing (Cat. No. 98TH8344), Vol. 2. IEEE, 877-880.
Andrea Cohen, Johannes L. Schönberger, Pablo Speciale, Torsten Sattler, Jan-Michael Frahm, and Mare Pollefeys. 2016. Indoor-Outdoor 3D Reconstruction Alignment. In Computer Vision ECCV 2016, Bastian Leibe, Jiri Matas, Nicu Sebe, and Max Welling (Eds.). Springer Interna- tional Publishing, Cham, 285-300.
Sukhendu Das, Arun Menon, and Koshy Varghese. 2017. Graph-based clustering for apictorial jigsaw puzzles of hand shredded content-less pages. Lecture Notes in Computer Science (includ- ing subseries Lecture Notes in Arlificial Intelligence and Lecture Notes in Bioinformatics) 10127 LNCS (2017), 135 147.
Prafulla Dhariwal and Alex Nichol. 2021. Diffusion Models Beat GANs on Image Synthesis. CoRR abs/2105.05233 (2021). arXiv:2105.05233
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. 2020. An Image is Worth 16x16 Words: Transformers for Image Recog- nition Scale. CoRR abs/2010.1 1929 (2020). arXiv:2010.1 19929 https://arzis: org/ aba/ 2010 11929
Herbert Freeman and L Garder. 1964. Apictorial jigsaw puzzles: The computer solution of problem in pattern recognition. IEEE Transactions on Electronic Computers (1964), 118-127.
David Goldberg, Christopher Malon, and Marshall Bern. 2002. A global approach to automatic solution of jigsaw puzzles. In Proceedings of the eighteenth annual symposium on Computational geometry. 27-87.
Jia Gong, Lin Geng Foo, Zhipeng Fan, Qiuhong Ke, Hossein Rahmani, and Jun Liu. 2023. Diff- Pose: Toward More Reliable 3D Pose Estimation. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition (CVPR).
Qipeng Guo, Xipeng Qiu, Pengfei Liu, Yunfan Shao, Xiangyang Xue, and Zheng Zhang. 2019. Star-transformer. arXiv preprint arXiv: 1902.019 (2019).
Zayd Hammoudeh and Chris Pollett, 2017. Clustering-Basee Based, Fully Automated Mixed-Bag Jigsaw Puzzle Solving. In Computer Analysis of Images and Patterns, Michael Felsberg, Anders Heyden, and Norbert Krüger (Eds.). Springer International Publishing, Cham, 205-217.
Peleg Harel and Ohad Ben-Shahar. 2020. Lazy caterer jigsaw puzzles: Models, properties, and a mechanical sastem-based solver. arXiv preprint arXiv:2008.07647 (2020).
Peleg Harel and Ohad Ben-Shahar. 2021. Crossing cuts polygonal puzzles: Models and Solvers. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 3084- 3093.
Jonathan Ho, Ajay Jain, and Pieter Abheel. 2020. Denoising Diffusion Probabilistic Models. arXiv preprint arxiv:2006..1 /239 (2020).
Daniel Hoff and Peter Olver. 2014a, Automatic Solution of Jigsaw Jouraal of Mathematical of Imaging and Vision 49 (05 2014).
11

---

Daniel J Hoff and Peter J Olver. 2013. Extensions of invariant signatures for object recognition.
Journal of mathematical imaging and vision 45, 2 (2013), 176-185.
Daniel Hoff and Peter J Olver. 2014b. Automatic solution of jigsaw puzzles. Journal of mathemat-
ical imaging and vision 49 (2014), 234-250.
Sepidehsadat Hosseini and Yasutaka Furukawa. 2022. Extreme Floorplan Reconstruction by
Cruuture-Hallucinatee Transformer Cascades. arXiv preprint arXiv:2206.006 (2022).
Bahjat Kawar, Michael Elad, Stefano Ermon, and Jiaming Song. 2022. Denoising Diffusion Restora- tion Models. In Advances in Neural Information Processing Systems.
Diederik P Kingma and Jimmy Ba. 2014. Adam: A method for stochastic optimization. arXiv preprint arXiv: arXiv:14/2.6900 14/2.6980 (2014).
Nikita Kitaev, Lukasz Kaiser, and Anselm Levskaya. 2020. Reformer: The Efficient Transformer.
In Interaational Conference on Learning Representations. https://opeareriew.nne wwt.wiiiiiiiee/orrrrr
1d TKggKKKRE
John Lambert, Yuguang Li, Ivaylo Boyadzhiev, Lambert Wixson, Manjunath Narayana, Will Hutchcroft, James Hays, Frank Dellaert, and Sing Bing Kang. 2022a. SALVe: Semantic Align- ment Verification for Floorplan Reconstruction from Sparse Panoramas. In ECCV.
John Lambert, Yuguang Li, Ivaylo Boyadzhiev, Lambert Wixson, Manjunath Narayana, Will Hutchcroft, James Hays, Frank Dellaert, and Sing Bing Kang. 2022b. SALVe: Semantic Align- ment Verification for Floorplan Reconstruction from Sparse Panoramas. In ECCV.
Canyu Le and Xin Li. 2019. JigsawNet: Shredded image reassembly using convolutional neural network and loop-based composition. IEEE Transactions Image Processing 28, 8 (2019), 4000-4015.
Haoying Li, Yifan Yang, Meng Chang, Huajun Feng, Zhihai Xu, Qi Li, and Yue-ting Chen. 2021b. SRDiff: Single Image Super-Resolution with Diffusion Probabilistic Models. CoRR Abs/2104.14951 (2021). arXiv:2104. 14951 https: http:///aaaavvooo/aa//////// wtwp/oo////......///////// 14951
Ru Li, Shuaicheng Liu, Guangfu Wang, Guanghui Liu, and Bing Zeng. 2021a. Jigsawgan: Auxiliary learning for solving jigsaw puzzles with generative adversarial networks. IEEE Transactions on Image Processing 31 (2021), 513-524.
Shiyang Li, Xiaoyong Jin, Yao Xuan, Xiyou Zhou, Wenhu Chen, Yu-Xiang Wang, and Xifeng Yan. 2019. Enhancing the locality and breaking the memory bottleneck of transformer on time series forecasting. Advances in neural information processing systems 32 (2019).
Xinghui Li, Kai Han, Shuda Li, and Victor Prisacariu. 2020. Dual-Resolution Correspondence Networks. In Conference on Neural Information Processing Systems (NeurIPS).
Cheng Lin, Changjian Li, and Wenping Wang. 2019. Floorplan-jigsaw: Jointly estimating scene layout and aligning partial scans. In Proceedings of the IEEE/CVF International Conference on Computer Vision. 5674-5663 5683.
Tsung-Yi Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. 2014. Microsoft COCO: Common Objects Context. Com- puter Vision ECCV 2014, David Fleet, Tomas Pajdla, Bernt Schiele, and Tinne Tuytelaars (Eds.). Springer International Publishing, Cham, 740-755.
Wen-Yan Lin, Siying Liu, Nianjuan Jiang, Minh N. Do, Ping Tan, and Jiangbo Lu, 2016. RepMatch: Robust Feature Matching and Pose for Reconstructing Modern Cities. In ECCV.
Ilya Loshchilov and Frank Hutter. 2017. Decoupled weight decay regularization. arXiv preprint arXiv:1711.5100 (2017).
William Marande and Gertraud Burger. 2007.
Mitochondrial DNA as Genomic Jigsaw Puz-
zle. Science 318, 5849 (2007), 415-415. wttp::///oo..gg/....1..... wwtp://///.......e....e. 1148803 148033
12

---

Smaragda Markaki and Costas Panagiotakis. 2022. Jigsaw puzzle solving techniques and applica- tions: survey. The Visual Computer (2022), 1-17.
Ricardo Martin Martin-Brualla, Yanling He, Bryan Russell, and Steven M Seitz. 2014. The 3d jigsaw puzzle: Mapping large indoor spaces. In European Conference on Computer Vision. Springer, 1-16.
Jonah C. McBride and Benjamin B. Kimia. 2003. Archaeological Fragment Reconstruction Using Curve-Matching. In 2003 Conference on Computer Vision and Puttern Recognition Workshop, Vol. 3-3. https://doo org/10...000. 1109/CVPRM 2003 10008
Chenlin Meng. Yutong He, Yang Song, Jiaming Song, Jiajun Wu, Jun Zhu, and Stefano Ermon. 2022. SDEdit: Guided Image Synthesis and Editing with Stochastic Differential Equations. In International Conference on Learning Representations.
Nelson Nauata, Sepidehsadat Hosseini, Kai-Hung Chang, Hang Chu, Chin- Yi Cheng, and Yasutaka Furukawa. 2021. GANu+: Generative Adversarial Layout Refinement Network towards Intelligent Computational Agent for Professional Architects. In Proceedings of the TEEF/CVF Conference on Computer Vision and Patlern Recognition. 13632-13641.
Alex Nichol and Prafulla Dhariwal. 2021a. Improved Denoising Diffusion Probabilistic Models. CoRR abs/2102.09672 (2021). arXiv:2102.099 https://arxii:
Alex Nichol, Prafulla Dhariwal, Aditya Ramesh, Pranav Shyam, Pamela Mishkin, Bob McGrew, Ilya Sutskever, and Mark Chen. 2021. GLIDE: Towards Photorealistic Image Generation and Editing with Guided Diffusion Models. CoRR abs/2112.10744 (2021). arX1v:2112.10741 10741
https://arxiv.or org/
Alexander Quinn Nichol and Prafulla Dhariwal. 2021b. Improved Denoising Diffusion Probabilistic Models. (2021). https://openreviien htw.:/nnwww......eeee/ooroouaiaa www.:eeeeeodddddeeeeeee NEXDKK8g2
Ture R Nielsen, Peter Drewsen, and Klaus Hansen. 2008. Solving jigsaw puzzles using image features. Pattern Recognition Letters 29, 14 (2008), 1924-1933.
Mehdi Noroozi and Paolo Favaro. 2016. Unsupervised learning of visual representations by solving jigsaw puzzles. In Computer Vision-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part VI. Springer, 69-84.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, et al. 2019. Pytorch: An imperative style, perf-mccformance deep learning library. Advances in neural information processing systems 32 (2019).
Mihai Pop. 2009. Genome assembly reborn: recent computational challenges. Briefings in Bioinformatics 10, 4 (05 2009), 354-366. https://dot. http://....rr//......///////.... arXiv:https:/academii http:////....cceocccc.....ocoobb
Gerald M Radack and Norman I Badler. 1982. Jigsaw puzzle matching using boundary-centered polar encoding. Computer Graphics and Image Processing 19, 1 (1982), 1-17.
Aditya Ramesh, Prafulla Dhariwal, Alex Nichol, Casey Chu, and Mark Chen. 2022. Hicrarchical text-conditional image generation with clip latents. arXiv preprint arXiv:2204.06125 (2022).
Robin Rombach, Andreas Blattmann, Dominik Lorenz, Patrick Esser, and Bjöm Om- mer. 2021. High- ReghhResolution Image Synthesis with Latent Diffusion Models. (2021).
LrXivccs.CV/2112.1077
Chitwan Saharia, Jonathan Ho, William Chan, Tim Salimans, David J Fleet, and Mohammad Norouzi. 2021. Image super-resolution via iterative refinement. arXiv:2104.0700 (2021).
Paul-Edouard Sarlin, Daniel DeTone, Tomasz Malisiewicz, and Andrew Rabinovich. 2020. Super- Glue: Learning Feature Matching with Graph Neural Networks. In CVPR. httpp://aaxit: org/ /tt.ab///991.11116.1
13

---

Hiroshi Sasaki, Chris G. Willcocks, and Toby P. Breckon. 2021. UNIT-DDPM: UNpaired Im-
age Translation with Denoising Diffussion Probabilistic Models. CoRR abs/2104.05358 (2021). arXiv:2104.05338
Mohammad Amin Shabani, Sepidehsadat Hosseini, and Yasutaka Furukawa. 2022. HiuseDiffu- sion: Vector Floorplan Generation via a Diffusion Model with Discrete and Continuous Denois- ing. arXiv preprint arXiv:2211.13287 (2022).
Mohammad Amin Shabani, Weilian Song, Makoto Odamaki, Hirochika Fujiki, and Yasutaka Fu- rukawa. 2021. Extreme Structure from Motion for Indoor Panoramas without Visual Overlaps. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV). https: //aminshabani github. wtt.:/...ooouulicaaiioioooo extreme httpr///ns...toeeeeeeee/e/////// stn/ppdtss/cccv2211 2088. pdf
Wenkang Shan, Zhenhua I Liu, Xinfeng Zhang, Zhao Wang, Kai Han, Shanshe Wang. Siwei Ma, and Wen Gao. 2023. Seaftusion-Basedd 3D Human Pose Estimation with Multi-Hypothesis Aggregation. arXiv preprint arXiv:2303. (2023).
Huang-Chin Shih and Chien-Liang Lu. 2018. Divide-and-conquer jigsaw puzzle solving. In 2018 IEEE Visual Communications and Image Processing (VCIP). IEEE, 1-2.
Noah Snavely, Steven M Seitz, and Richard Szeliski. 2006. Photo tourism: exploring photo collec- tions in 3D. ACM siggraph (2006), 335-846.
Yang Song, Jascha Sohl-Dickstein, Diederik P. Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. 2020. Score- Baoee- Generative Modeling through Stochastic Differential Equations. CoRR abs/2011.13456 (2020). arXiv:2011. 1. 13456 https: /ttps://arxiv. wtwp://.rrrbab////./// 13456
Jiaming Sun, Zehong Shen, Yuang Wang, Hujun Bao, and Xiaowei Zhou. 2021. LoFTR: Detector- Free Local Feature Matching with Transformers. CVPR (2021).
Corey Toler-Franklin, Benedict Brown, Tim Weyrich, Thomas Funkhouser, and Szymon Rusinkiewiez. 2010. Multi-feature matching of fresco fragments. ACM Transactions on Graphics (TOG) 29, (2010), 1-12.
Qiuhong Anna Wei, Sijie Ding, Jeong Joon Park, Rahul Sajnani, Adrien Poulenard, Srinath Sridhar, and Leonidas Guibas. 2023, LEGO-Net: Learning Regular Rearrangements of Objects in Rooms. arXiv preprint arXiv:2301.09629 (2023).
Haim Wolfson, Edith Schonberg, Alan Kalvin, and Yehezkel Lamdan. 1988. Solving jigsaw puzzles by computer. Annals of Operations Research 12, 1 (1988), 51-64.
Julia Wolleb, Robin Sandkühler, Florentin Bieder, Philippe Valmaggia, and Philippe C. Cattin. 2021. Diffusson Models for Implicit Image Segmentation Ensembles. (2021).
Wenming Wu, Xiao-Ming Fu, Rui Tang, Yuhan Wang, Yu-Hao Qi, and Ligang Liu. 2019. Data- driven interior plan generation for residential buildings. ACM Transactions on Graphics (TOG) 38, 6 (2019), 1-12.
Zhenpei Yang, Jeffrey Pan, Linjie Luo, Xiaowei Zhou, Kristen Grauman, and Qixing Huang. 2019. Extreme Relative Pose Estimation for RGB-D Scans via Scene Completion. In Proceedings of the IEFF/CVF Conference on Computer Vision and Pattern Recognition (CVPR).
Zhenpei Yang. Siming Yan, and Qixing Huang. 2020. Extreme relative pose network under hybrid representations. Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition. 2455-2464.
Kwang Moo Yi, Eduard Trulls, Yuki Ono, Vincent Lepetit, Mathieu Salzmann, and Pascal Fua. 2018. Learning to find good correspondences. In Proceedings the IEEE conference on computer vision and pattern recognition. 2666-2674.
Min Zhao, Fan Bao, Chongxuan Li, and Jun Zhu. 2022. Egsde: Unpaired image-to-image translation via energy-guide stochastic differential equations. arXiv preprint arXiv:2207.0663: (2022).
14

---

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-shariin formers 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 out- put 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.
among input tokens. These two features make Trans-
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 ro- tations as explained below. An input room layout is represented as a 20-channel 256 x 256 semantic segmentation image, where there are 20 room/door types. The room center aligned with the center of an image. An output room position is represented as a 256 x 256 room occupancy image, which 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. We use VisionTrransfommer Dosovitskiy et al (2121) with CNN de- coder 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 hown-ssapling blocks, converting each input room layout to feature map of dimension 512. Each feature map (correspond- ing to a room layout) will become one input token for the Transformer. Input sequences length is equal to the number of rooms in 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 x 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 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 O. We flatten the conditions per room and then use 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, linear layer converts 256D output to 4D (i.e., for the position and 2 for the rotation).
Shabani el al. Shabani et all (2021I) takes the input layout of each room with the resolution of 256 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 Shabam et (20221) 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 significantly larger than the one in Shabani et all (2022I,, enabling us to randomly select a positive or a negative candidate in each iteration and therefore remove the class imbalance weight used in Shabam el all (202TT) During the training for each house, we randomly select a GT with the label I or faulty candidate with the label [O, 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.
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.
15

---

Harel et al. Harel and Ben-Shahar (202T) proposed a two-step algorithm for CJP. Their method con- siders two types of constraints to find plausible mates based 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 multi-body spring-mass system. We utilize the authors' provided implementation for comparison with our method. With 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 metries. 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
4
7
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 specific room type based the total number of rooms in the house.
We normalize the puzzles/floorplass for each task and dataset by scaling them to fit within a 1 x 1 square, and we also resize all corresponding images to dimensions of 256 x 256 in the case of pic- torial CJP. While this normalization process does not introduce any essential additional information during testing in CJP and VJP, could potentially enables the network to cheat in RLA, as the longer extent of arranged floorplans is always fixed to 1. To address this issuc, during testing in RLA, apply a random scaling factor in the range of [0.8, 1.0] to the room shapes each house. In the subsequent sections, provide 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 Ik testing puzzles. These puzzles were created by randomly selecting 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, 6,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 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 Shabar (2021) method which generate a convex polygon and cuts by 3 to lines. There are 1719, 6046, 15854, 14521, 6905, 10929, 12065. 8361, 6521, 8192, 7663, 4642, 1508, 73, and puzzles with 3, to 18 pieces respectively. Each piece has minimum,
https//eww..s..guu.....uu
16

---

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 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 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 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 corre- sponding full dataset, consisting of houses with at most rooms. The small datasets are created for Shabani ef al., which is not scalable many rooms. Our method stochastic and shows both the mean and the standard deviation.
Dataset
Small RPLAN
Full RPLAN
Small MagicPlan
Full Jigsaw Plan
Metric
MPE (1)
GED (4)
MPE (L)
GED (1)
MPE (1)
GED (1)
MPE (L) GED (1)
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
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+08
1.0+0.4
27.9+0.7
2.7+0.5
Figure II 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 [4] shows five pose estimation results by our system while varying initial noise ST. While there are minor differences, room arrangements are similar and close to the ground-truth, indicating that the Diffusion model capable of producing consistent results given enough constraints as pose estimation system, as opposed to a generative model whose original goal is to create a diverse set of answers.
overall
Input
== =000
t=695
207
E
Ground Truth
Figure 9: Visualization of predicted layouts step "t"s. At 1=100, position parameters at each corner are initialized by Gaussian noise, and at t=0, there the final predicted layout. The top row shows the predicted layout without averaging/voting, and the bottom row shows with averag- ing/voting. To make it more clear, we show doors by their corners.
17

---

Input
Shabani et al.
TransRaster
Vector
Ours
Ground Truth
Figure 10: Qualitative evaluations of our approach against the three competing methods. The top two rows are from Small MagicPlan. The bottom two rows 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 the MagicPlan dataset in case of layout rangment task. This part of supplementary will present the same study results on RPLAN dataset. Table Ø shows the impact of our attention module and door matching loss on performance and Ta- ble & shows the impact of noise in the room type and door detection on our performance, although there is performance drop, our method still works better than the competing methods.
Table 7:
Co attention mechanisms (P-SA, G- SA) and the door matching loss (Lmatch) Full RPLAN is used. / /indicates the feature be- ing used. In case match "Doors" means matching loss has been applied only door corners and "All corners" means matching loss has been applied all corners including door corners.
Table 8: Effects of the room-type (R-type) and the Door information. Full RPLAN used. v indicates the information being used. When room-type
used, set zero vector as room-type one-hot vector. When the door in- formation is not used, we do not pass the door- corner nodes to the network.
Train
Test
P-SA
G-SA
Tmmatch
MPE (L)
GED (J)
MPE([)
GED (L)
R-Type
Door
R- R-Type
Door
17.3
24.2
1.5
16.4
1.5
Doors
36.9
2.4
f
15.1
1.9
Doors
22.1
1.1
/
14.3
1.9
All Corners
10.7
0.9
10.5
0.9
Doors
10.5
0.9
18

---

Prodicted Atrangement
Ground Truth
Predicted Arrangement
Ground Truth
Predicted Arrangement
Ground Truth
Predicted Arragementt
Ground
Predicted Arrangement
Predicted Arrangement
Ground
12T
Arrangemett
Ground
Predicted Arrangement
Predicted Arrangement
Prodicted Arrangement
Ground Truth
Prodicted Arrangement
Ground Truth
Predicted Arrangement
Ground Truth
Qualitative evaluations of method for Full MagicPlan dataset without GT rotations top two rows, and Full RPLAN dataset without GT rotations bottom two rows. We show edges and corners here show overlaps and noisy annotations more clear.
Figure 11:
C.2 Additional ablation studies on puzzle solving
We have provided additional qualitative results of our method in Figure L and Figure IGA including noisy samples or samples with missing duplicate pieces. In case of missing and duplicate exper- iment, we repeat (remove) each piece with probability of 10%. Table El presents the evaluation metrics of missing and duplicate experiments.
C.2.1 Pictorial Cross-cut Jigsaw Puzzle
To enhance the integration of image information into our pictorial puzzle diffusson models, employed a two-step approach. Firstly, we pretrained an auto-encoded utilizing the puzzle pieces. This auto-encoder served the image embedder for our diffusion model, enabling the conversion of each puzzle into a compact 128D feature vector.
The pretraining process involved training the model to downsample an input image of dimensions 3 x 256 256 compressed representation of size 32 x 2 x 2 within the encoder, and subse- quently reconstructing the original image size in the decoder. We employed the mean squared error
19

---

Input
Before Averaging
After Averaging
Ground Truth
Prediaed Anveporant
Groaae Frars
Fredioud Artrppatee:
Gronad Tours
RLA arrangement results with Figure 13: The final room arrangement before and af- Full MagiePlan dataset. On left two suc- ter averaging. Our diffusion model estimates a room cessful cases, on right two failed cases,. Our position/rotation at each room corner, which may not failures are often attributed to 1) Rare build- be consistent in a room. The final arrangement ing architecture (top-right) and 2) Inherent obtained by taking the average position and rotation ambiguity (bottom-right ), whose tasks are within each room. challenging for humans.
Figure 12:
Input
Predicted Arrangements
Ground Truth
Figure 14: diffusion model is stochastic and produces different result every time. The middle rows show five different pose estimation results. The top (resp. bottom) is from Full RPLAN (resp. Full MagicPlan) dataset.
(MSE) loss function during training. However, to focus our model's attention on learning the texture features, given that the diffusion model already captured the geometry features, applied the loss function exclusively to the pixels within the puzzle piece.
By adopting this selective application of the loss function, prioritize the acquisition of texture- based details, as the geometric characteristics are already embedded within the diffusion model.
Quantitatively, we also evaluated our method on the full Cross-cut dataset to measure the effective- ness of the pictorial information compared apictorial scenario. We found that the model converges faster when using pictorial information while it achieves slightly better overlap score of 0.9417 com- pared 0.9398 apictorial scenario. Figure shows additional qualitative results our method compared to Harel ef al. Harel and Ben-Shabar (2021))
20

---

Table 9: Effects the Missing or Duplicate pieces in puzzle solving problem. inddccates missing duplicate piece were presented during test time. In training time do not have duplicate or missing piece presented to show our model robustness to unseen noise during test.
Cross-cut
Voronoi
Missing Duplicate Overlap (t)
Precision (t)
Recall (f)
Overlap (t)
Precision (t)
Recall (t)
0.88
0.92
0.82
0.68
0.68
0.56
0.92
0.97
0.88
0.67
0.71
0.57
0.94
0.97
0.91
0.70
0.78
0.60
21

---

Ours
GT
Ours
GT
Noise Level 2
Noise Level =
Ours
GT
Ours
GT
Duplicate
Missing
Ours
GT
Ours
GT
Noise Level =
Noise Level
Ours
GT
Ours
GT
Duplicate
Missing
Ours
GT
Ours
GT
Noise Level 2
Noise Level =
Ours
GT
Ours
GT
Duplicate
Missing
Figure 15; Additional qualitative results of Voronoi jigsaw puzzle presented in four different setups: 1) No noise, 2) Noise level 2, 3) Missing piece, and 4) Duplicate piece (D indicates the duplicated pieces).
22

---

Ours
GT
Ours
GT
Noise Level = 2
Noise Level
Ours
GT
Ours
GT
Duplicate
Missing
Ours
GTT
Ours
GT
Noise Level =
Noise Level
Ours
GT
Ours
GT
Missing
Duplicate
Ours
GT
Ours
GT
Noise Level 0
Noise Level
Ours
GT
Ours
GT
Missing
Duplicate
Figure 16: Additional qualitative results of Cross-cut jigsaw puzzle are presented in four different
setups: 1) No noise, 2) Noise level 2, 3) Missing piece, and 4) Duplicate piece (D indicates the duplicated pieces).
23

---

Harel et al.
Ours
GT
Figure 17: Additional qualitative results of pictorial Cross-cut jigsaw puzzle compared to Harel et al. HHaell and Ben-Shahad (202T).
24


Reviews:
[{"review_id": "ewTklJiA20", "reviewer": "Reviewer_c71V", "summary": "This paper introduces a diffusion model for solving 3 different spatial puzzle tasks: cross-cut jigsaw, voronoi jigsaw, and room layout arrangement. Their method achieves SOTA results while being much faster than previous works, allowing them to handle larger puzzles than previous methods could. They also demonstrate greater robustness to noisy data inputs compared to previous methods, despite not being trained on noisy data. Ablations show the contributions of proposed losses and network components. They also present two new datasets: a synthetic one for the Voronoi puzzle task, and a real room layout one from MagicPlan.", "questions": "- As shown in Fig. 7 of the supplementary PDF, the method produces slightly different results due to stochasticity; how many samples were generated when computing the results in Tables 1 and 2?\n- L243 says that TransVector achieves a better MPE score than Shabani et al, but that doesn’t appear to be the case in Table 1?\n- In Table 2, are the CJP and VJP results obtained using the same model, or is a separate model trained for each task?\n- For Table 1, separate models are trained for MagicPlan and RPLAN, but does a model trained on MagicPlan exhibit some generalization to RPLAN and vice versa?", "limitations": "Limitations were discussed adequately.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 4, "contribution": 3, "strengths": "- The paper is overall well-written and easy to understand.\n- In terms of originality, this paper is the first to tackle these kinds of layout arrangement problems using diffusion models.\n- The proposed method is sound and clearly described.\n- The method is a significant improvement over existing methods while using a very different approach.", "weaknesses": "- No major weaknesses\n- At zero noise level (Figs 5 and 6), the proposed method seems to be less precise at aligning the pieces compared to Harel et al. (while the general layout is correct, there are small alignment errors). Of course, Harel et al. fails for more difficult cases, but for these easier cases, it seems slightly more accurate than the proposed method. (I am however not very familiar with this task so I am not sure how important these small inaccuracies are in practice.)\n- Some details are unclear (see “Questions”)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your positive, insightful and valuable comments and suggestions which are very crucial for improving the quality of our manuscript.\n\n\n---\n**1. At zero noise level (Figs 5 and 6), the proposed method seems to be less precise at aligning the pieces compared to [6].**\n>We agree that Harel et al. achieves slightly more precise alignment through their heuristic method in the case of zero noise. \nHowever, this is not an issue. We conducted an experiment (see the attached pdf Fig. 2) showing that our approach easily achieves equally precise alignment by a simple post-processing heuristic, which is a minor variant of a loop merging process of an existing work [1]. Our method is also noticeably faster than [6]. We will add the details of the post-processing step and more experimental results to the paper.\n\n>[1] Chen, Jiacheng, et al. \"Floor-sp: Inverse cad for floorplans by sequential room-wise shortest path.\" Proceedings of the IEEE/CVF International Conference on Computer Vision. 2019.)\n\n---\n\n**2. How many samples were generated when computing the results in Tables 1 and 2?**\n\n>We thank for the question and will clarify in the paper. We ran our system 5 times and reported the mean.\n\n---\n\n**3. Typo on line 243.**\n\n>We apologize for the oversight and thank you for bringing that to our attention. That sentence is indeed a typo, and we will drop the phrase from the text.\n___\n\n**4. Are the CJP and VJP results obtained using the same model, or is a separate model trained for each task?**\n\n>We thank for the question and will clarify in the paper. The models are independent and trained separately.\n\n---\n**5. Does a model trained on MagicPlan exhibit some generalization to RPLAN and vice versa?**\n\n>We thank for a great question. We conducted additional experiments, which will be added to the paper. Concretely, we trained our model with MagicPlan and tested it with RPLAN. (MPE/GED) scores are (15.62/1.90), compared to (10.55/0.97) when trained and tested with RPLAN. Similarly, (MPE/GED) scores are (48.48/4.68) when trained with RPLAN and tested with MagicPlan, compared to (40.81 and 3.09) when trained and tested with MagicPlan. The performance drops are understandable, given that Magicplan encompasses real-world data contributed by users, while RPLAN contains synthetic noise-free floorplans, designed by professional architects. It is worth noting that our cross-dataset results are still noticeably better than in-dataset results by the TransVector baseline."}, {"review_id": "oUTp4AKfuo", "reviewer": "Reviewer_qFVS", "summary": "The paper presents a diffusion based method to tackle jigsaw puzzle solving task. This task has applications in artwork restoration, room layout estimation, etc. The paper also introduces a room layout and arrangement dataset. The authors compare their work with previous state of the art approaches and a transformer based baseline, and outperforms all these methods. ", "questions": "### Copying my comments from weakness section\n\n1.  Authors state that redundant representation helps their system, and have shown one ablation for it. However, it would be interesting to see an ablation that shows how the method performs when the the position and orientation estimation of the piece are not concatenated with all the corners but passed as separate input.\n\n2. Some discrepancy between figure and supplementary video: In fig 2, it seems like that random noise is added to each corner independently and that noise changes the change of each piece. In the supplementary video, however, it seems that same random noise is added to all the corners of each piece, which results in shape of each piece staying consistent. So which one is it?\n\n3. If random noise is added to each corner independently for each piece, how do authors maintain the ordering of the corners when passing it to transformers? Won’t de-noising process change the ordering of corners? \n", "limitations": "Yes, authors have addressed the limitations well", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The paper presents an interesting approach to tackle jigsaw puzzle solving problem using diffusion models. The corner coordinates and rotation of puzzle pieces are used as input, and diffusion model is tasked with predicting the noise in those inputs. \n2. Unlike prior methods, this method is not limited to simple puzzles or require pairwise comparison between puzzle pieces, which makes it more efficient.\n3. The method achieves better performance compared to previous state of the art methods and a transformer based baseline that directly predicts the output instead of doing iterative denoising.\n4. Ablations show that all new modules/loses introduced in the paper contribute towards the final performance.", "weaknesses": "### Paper clarity comments\n\n1. Section 3 is a bit hard to follow. \n2. What does line 140 mean: “position/rotation of the rth room/piece stored at the ith corner”? Does it mean the piece that has corner i as one of the corners?\n\n### Missing ablation\n3. Authors state that redundant representation helps their system, and have shown one ablation for it. However, it would be interesting to see an ablation that shows how the method performs when the the position and orientation estimation of the piece are not concatenated with all the corners but passed as separate input.\n\n### Discrepancy between paper and supplementary\n4. Some discrepancy between figure and supplementary video: In fig 2, it seems like that random noise is added to each corner independently and that noise changes the change of each piece. In the supplementary video, however, it seems that same random noise is added to all the corners of each piece, which results in shape of each piece staying consistent. So which one is it?\n5. If random noise is added to each corner independently for each piece, how do authors maintain the ordering of the corners when passing it to transformers? Won’t de-noising process change the ordering of corners? ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your positive, insightful and valuable comments and suggestions which are very crucial for improving the quality of our manuscript.\n\n--- \n\n**1. Section 3 is a bit hard to follow**\n\n>We thank for the comment and will do our best to clarify our writing. We would appreciate any further specific comments or suggestions.\n--- \n\n**2. What does line 140 mean: ''position/rotation of the rth room/piece stored at the ith corner''? Does it mean the piece that has corner i as one of the corners?**\n\n> We thank for the question and will clarify in the paper. ($r$) denotes the index of a room/piece in a puzzle. For example. ($r$) is either 1, 2, or 3 for a puzzle of 3 pieces. Similarly, ($i$) denotes the index of a corner in a room/piece. Therefore ($C^r_i$) denotes the i-th corner of the r-th room/piece. Note that our approach infers the position/rotation information at every corner of a room/piece and uses their average to determine the final position/rotation of the room/piece.\n\n--- \n\n**3. it would be interesting to see an ablation that shows how the method performs when the the position and orientation estimation of the piece are not concatenated with all the corners but passed as separate input.**\n\n>That is indeed an interesting experiment. In order to study network performance while considering position and orientation as separate inputs, we conducted two experiments. In the first experiment, we utilized a single input token per piece to predict both piece position and rotation of that piece, while incorporating corners information as separate tokens. In the second set of experiments, we employed one input token for each corner to predict the position and rotation of each piece, with corner information again represented as distinct tokens. We applied both of these methods to the MagicPlan dataset.\n>\n>For the first setup, we achieved MPE and GED values of 49.91 and 4.96, respectively. Similarly, for the second setup, the corresponding values were 42.45 for MPE and 3.22 for GED. The performance drop is significant in the first setup, and although the second setup yields more comparable results, it still demonstrates lower performance when compared to our setup (MPE/GED of 40.81/3.09). Additionally, it involves twice the number of input tokens compared to our setup, resulting in higher computational costs. We will provide details of this experiment in the paper.\n \n\n---\n\n \n\n**4. It seems that same random noise is added to all the corners of each piece.**\n\n>Different noises (but from the same Gaussian distribution) are added to different corners, with each corner contributing to a prediction for the position or rotation of the corresponding piece. The final prediction is derived by averaging the predictions from all corners of a piece. Nevertheless, when visualizing the diffusion process, there are two options: utilizing the averaged output of the entire piece for corner positions, or solely relying on the output of each corner itself. Figure 2 in the supplementary material and the video from 00:57 to 1:26 effectively illustrate this concept.\n\n___\n\n**5. How do authors maintain the ordering of the corners when passing it to transformers?**\n\n>We thank for the question. The feature embedding contains the corner index information (the second term of Eq. 3 in Section 4.2), which allows our system to keep track of the order of corners. We will clarify our explanation at Lines 156-158 of the main paper."}, {"review_id": "VvOLgLkgew", "reviewer": "Reviewer_AMsE", "summary": "This paper investigates the use of diffusion models to solve jigsaw puzzles. Puzzle pieces are model as a sequence of corners and the diffusion process consists in adding noise/denoising the position and orientation of the corners, conditionally to the original piece shapes. To have fragments snap into place, an additional loss for corner matching is added. The diffusion process is test on synthetic puzzles as well as on a floorplan dataset and shows that the proposed method in improving over the existing literature.", "questions": "- Is it possible to solve the regular n x n jigsaw puzzle (for example by assigning to each fragment its closest fixed position) and compare to regular method that solve the combinatorial problem (like [35,36] and similar methods) using fragment position accuracy and puzzle solving accuracy? This would allow for much more interpretable results.\n- What are the score for pictorial CJP?", "limitations": "The paper addresses limitations in the conclusion by mentioning the requirement of large scale training sets.", "rating": 6, "confidence": 2, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Given the combinatorial nature of jigsaw puzzles, using a denoising diffusion process, which runs in a fixed number of steps, is a promising idea. It opens the door of using diffusion process to solve discrete combinatorial tasks, which is an area where learning based method are still struggling to dominate (setting aside Deep MCTS). The execution in this paper is well done, the additional loss for corner matching is sound. The experiments are OK. Code was provided.", "weaknesses": "Not much!\n- One complaint I have is that the evaluation metrics are not very interpretable. The MPE depends on the size of the puzzle and the pieces (a 2 pix error is not the same on a 5 pix fragment and on a 150 pix fragment), GED is even worse and precision/recall tend to saturate on the proposed datasets. Also, pictorial CJP have no quantitative results mentioned.\n- I was having the impression that the literature on jigsaw puzzle solving using deep learning is a bit sparse ([26, 35, 36]), given the number of relevant papers that a google scholar search returns when querying \"deep learning jigsaw puzzle\". More comparison with how deep learning was used to solve the problem would have be nice to highlight the radical difference this approach proposes.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for your positive, insightful, and valuable comments.\n\n---\n**1. Interpretability of the metrics.**\n\n>We thank the reviewer for pointing this out. While these metrics might not be deemed ideal, it worth to mention that for layout arrangement, as every pixel is equivalent to a real world distance (measured in meters/feet), which can be determined using the room layout scale of the house, MPE can easily converted to real-world distance, as highlighted by Shabani et al. [20]. In addition, to further address the reviewer's concern, we have devised a new **\"Weighted MPE\" (WMPE)** metric, which normalizes the error per room/piece by the square root of the room/piece area:\n>$\n\\text{WMPE} =\\frac{1}{N} \\sum_{i=0}^{N} \\frac{\\text{Positional Error of } x_{i}}{\\sqrt{\\text{ Area of }x_{i} }}\n$\n>\n>The WMPE metric scores are 0.109 (ours) and 0.451 (TransVector) for the full RPLAN dataset, and 0.341 (ours) and 0.564 (TransVector) for the full MagicPlan dataset. We will add the results to the paper.\n\n>The GED metric was borrowed from Nauata et al. [56] to assess room connectivity, which is established through doors/frames, stands as a pivotal factor in house design. The precision/recall metric has been the standard in the recent jigsaw puzzle literature [6,26].\n\n---\n**2. Quantitative results of CJP.**\n\n>Due to the time complexity of Harel et al. especially in the case of pictorial puzzles, they did not report any number for their method on pictorial data. For the same reason, in order to compare Harel et al. to ours, we calculated the metrics on 20 samples for pictorial CJP. Our method achieves 0.953, 0.978, and 0.930, respectively for, IoU, Precision, and Recall, compared to 0.935, 0.934, and 0.947 of Harel et al. Thank you for pointing out, we will add the details to the paper.\n\n\n---\n**3. Additional related works on jigsaw puzzles.**\n\n>We thank for the suggestion. We will add and discuss more references on deep-learning based jigsaw puzzle papers. However, the majority of these papers tackle square puzzles and cannot be compared on our tasks [25]. Our work pursues a broader and more adaptable approach for more challenging tasks, and needs to compare against a few similarly adaptable approaches (e.g., [Harel et al.] or [Shabani et al.]).\n\n---\n**4. Is it possible to solve the regular n x n jigsaw puzzle?**\n\n>Yes, although the focus of our paper is mainly on the geometry side of the puzzles, our method can be used for pictorial square puzzles as well. In this regard, we did an experiment on solving 3x3 pictorial square puzzles with similar implementation details as the pictorial CJP. Our model achieves the following metrics:\n>- 90.98% in *Direct Comparison*, which denotes the portion of puzzle pieces within the reassembled puzzle that have been positioned correctly,\n>- 88.87% in *Neighbour* score, meaning the percentage of paired neighboring pieces that are correct,\n>- 74.54% in *Perfect Reconstruction*, which measures the percentage of perfectly reassembled puzzles.\n>\n>We need to emphasize that our approach tackles puzzles with randomly initialized pieces in a continuous space. In contrast, prior methods such as [35,36] focus solely on a subset of potential discrete permutations. For instance, Table 4 in reference [35] displays a maximum of 1000 permutations selected from the entire pool of 9! (~363k) possible permutations. Even with this reduced set, the challenge posed by the number of permutations persists, especially compared to the results from consideration of a mere 10 or 100 permutations. Few qualitative samples are provided in the attached PDF file. We will add this to the main paper."}, {"review_id": "Dx78NWu6gA", "reviewer": "Reviewer_TNRQ", "summary": "This paper introduces a novel approach to puzzle solving and room floorplan arrangement by utilizing a conditional generation process based on the denoising diffusion model. The model effectively reconstructs the original polygonal coordinates, representing the spatial arrangement, during the reverse diffusion process. Notably, the proposed method demonstrates robustness to data noises, which enhances its practical applicability.\n\nIn order to train the diffusion model, the authors introduce two new datasets: a synthetic jigsaw dataset and a real floorplan dataset with room layout pieces.\n\nThe experimental results showcase outstanding performance in both puzzle solving and room floorplan arrangement tasks.", "questions": "1. Can the proposed method handle puzzles that are cut into squares?\n2. Can a baseline be composed of TransVector-based approaches with multiple iterations of optimization?\n3. Could you provide more details on the \"Averaging/voting\" process in Figure 2 of the supplement? \n4. How does the computational cost of the proposed method compare to existing methods?", "limitations": "N/A", "rating": 7, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The utilization of positional information as the signal in the diffusion process is a relatively novel idea.\n2. Introduces an interesting approach by framing spatial arrangement, puzzle solving and floorplan registration tasks as a conditional generative process.", "weaknesses": "1. The paper exhibits a writing problem where the description of the jigsaw solving and layout arrangement tasks lacks strong connections.\n2. It would be beneficial to provide more detailed explanations about the jigsaw part directly in the main paper.\n3. The training of the diffusion model requires a large amount of data.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for your positive, insightful and valuable comments and suggestions which are very crucial for improving the quality of our manuscript.\n\n---\n**1. Additional details of the jigsaw part in the main paper and lack of connection between layout arrangement and jigsaw puzzles.**\n\n>We thank for the comment and will do our best to clarify the details of the layout arrangement and jigsaw puzzle tasks. Both tasks assemble scattered geometric pieces into a single coherent shape, an image, or a floorplan of a functional house. Careful spatial reasoning is the key to the success of a method. We will clarify this in the paper and would appreciate any further specific feedback/comments on the clarity.\n\n---\n\n**2. The training of the diffusion model requires a large amount of data.**\n\n>We agree with the reviewer, where one of our key contributions is the introduction of a new large-scale real-world dataset, MagicPlan, which helps us train our diffusion-based model. This contribution is precious given the scarcity of substantial real estate data, owing to privacy and licensing constraints. In addition, this challenge does not apply to Jigsaw puzzle tasks whose datasets are often synthetic and scale up easily.\n\n--- \n\n**3. Can the proposed method handle puzzles that are cut into squares?**\n\n>We thank for the question and have conducted experiments to show that our method can solve square jigsaw puzzles.\n>\n>Specifically, we have solved 3x3 pictorial square puzzles with a similar implementation as for the pictorial CJP. Our model achieves the following numbers:\n>- 90.98% in *Direct Comparison*, which denotes the portion of puzzle pieces within the reassembled puzzle that have been positioned correctly,\n>- 88.87% in *Neighbour* score, meaning the percentage of paired neighboring pieces that are correct,\n>- 74.54% in *Perfect Reconstruction*, which measures the percentage of perfectly reassembled puzzles.\n>\n>We want to emphasize that our approach randomly initializes and solves the arrangements in a continuous space. In contrast, prior methods such as [35,36] focus solely on a subset of potential discrete permutations. For instance, Table 4 in reference [35] displays a maximum of 1000 permutations selected from the entire pool of 9! (~363k) possible permutations. Even with this reduced set, the challenge posed by the number of permutations persists, especially compared to the results from consideration of a mere 10 or 100 permutations. A few qualitative samples are provided in the attached pdf file. We will add this to the main paper.\n\n___\n\n**4. Can a baseline be composed of TransVector-based approaches with multiple iterations of optimization?**\n \n>We thank for the suggestion and have implemented this new baseline as an additional experiment. Drawing inspiration from the iterative approach employed by Housegan++ [56] during training, we integrated a 50% probability for each room to remain fixed. In cases where rooms were designated as fixed, we input the ground truth position (GT position) to the network, enabling the network to learn the task of arranging non-fixed rooms by leveraging information from the fixed ones. \n>\n>During the testing phase, previously predicted locations were introduced to the network with a 50% probability per room, effectively offering potential input constraints and facilitating iterative design refinement. This process was repeated a total of 10 times. Consequently, the performance improved, yielding MPE of 39.12 and GED of 2.18 on the full RPLAN dataset, while the full MagicPlan dataset reflected 47.85 and 5.49, respectively for MPE and GED. Despite noteworthy improvements, the results remained significantly inferior to our diffusion-based approach. We will add this to the paper.\n\n---\n\n**5. Could you provide more details on the \"Averaging/voting\" process in Figure 2 of the supplement?** \n\n>We thank for the question and will clarify the following details in the paper. For puzzle solving, the determination of the location and rotation of each piece involves computing the average location and rotation of its corners. However, in the context of room layout arrangement, where rotations adhere exclusively to the Manhattan directions (0, 90, 180, and 270 degrees), a voting system for rotation is employed. This includes assigning each corner a vote, aligned with the rotation closest to its predicted value. The final rotation of the piece corresponds to the angle that accumulates the highest number of votes. For position, we use a similar approach to puzzle solving. Predicted positions for all corners within a room are averaged to establish the position of the room within the floorplan.\n>\n>When visualizing the diffusion process, there are two options: utilizing the averaged output of the entire piece for corner positions, or solely relying on the output of each corner itself. Figure 2 in the supplementary material and the video from 00:57 to 1:26 provides both visualizations.\n\n---\n**6. How does the computational cost of the proposed method compare to existing methods?**\n\n>The diffusion model and TransVector implementations each utilize approximately 4 million parameters. On the other hand, Harrel operates as a non-learning-based method, devoid of any parameters to be learned. In the case of Shabani et al. [20], their approach encompasses a blend of a heuristic technique for generating potential candidates alongside a deep learning model for scoring said candidates. While the model itself incorporates a lesser number of parameters, the heuristic component demands notably more time due to the generation of all possible candidates via door connections. The training and the inference time of our method and the competing ones are given at L183-187 and L234-238 of the paper."}, {"review_id": "Dfw5eEmrLh", "reviewer": "Reviewer_zazU", "summary": "This paper applies conditional diffusion model to address the problem of spatial puzzle solving and show its real application on Cross-cut Jigsaw Puzzle (CJP), Voronoi Jigsaw Puzzle (VJP) and room layout arrangement (RLA). Unlike previous methods that rely on enumerating and verifying pairwise alignment and may struggle when the complexity of global arrangement increases, this approach aligns all pieces in one pass. Qualitative and quantitative evaluations show the proposed method outperforms competitive methods. ", "questions": "+ Experiments:\n    + It seems better to derive more analysis on ablation studies. For example, in Tab. 3 lines 5 and 8, why does matching loss applying on all corners have a slightly lower performance than the optimal choice?  Besides, feature redundancy at each corner seems critical, what is the intuition behind this design? Can it be explained by providing more training examples to the model?\n    + There should be more experimental results. Does reference [21] have more competitive results on Room layout arrangement? It is better to report the results of [21] in Tab. 1.\n    +  The evaluation scheme is not clear, since the ideal result is deterministic, how does the method cope with the variability of the DDPM output during the evaluation?\n\n+ Typo: Tab. 2 Caption: quantitative\n", "limitations": "The authors discuss some limitations about the performance and demands on the data. I would like to see more discussions about the quality of the results, which can be more inspiring. ", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ The writing is easy to follow.\n+ The proposed method verifies the potential of the generative model on spatial puzzle solving, which may inspire to address more complex spatial puzzles.\n+ Experiments with noisy spatial puzzles demonstrate the robustness of the proposed methods.", "weaknesses": "+ The application and the method are not properly motivated. I cannot come up with a case that we need to put together multiple floorplan regions. It looks a little weird. Besides, the other two applications, i.e., CJP and VJP, are too simple. \n\n+ I did not see how the method solves complex cases as raised in line 30.  Full MagicPlan (RPLAN) has at most 10 (8) room patches, while Small MaginPlan (RPLAN) has at most 6 (6) room patches, thus they have minor differences in complexity. The search space of these examples is small and may not need a learning-based method.  \n\n+ Methodology: The canonical diffusion denoising probabilistic model requires quite some denoising steps, the proposed method does not apply any modifications to remedy the low time efficiency.\n", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "\nWe thank the reviewer for your positive, insightful and valuable comments and suggestions which are very crucial for improving the quality of our manuscript.\n\n---\n\n**1. Applications of the room layout arrangement and hardness of the jigsaw puzzle solving.**\n\n>The room layout arrangement is an emerging problem for the real estate industry, with big companies like Ricoh, Zillow, and MagicPlan actively pursuing solutions. For instance, there exist mobile applications that capture individual room layouts, necessitating manual integration to produce the full house floorplans. Our method is a step forward toward automating this process. Please also refer to L11-L13, L22-L24, L106-109, and L201-207 of the paper, mentioning the practical impact of the proposed research.\n > \n>Moreover, tackling the CJP and VJP demands at least several minutes of human effort. In contrast, our method accomplishes this task within mere seconds. We believe extending our method to similar tasks and datasets could assist in addressing formidable challenges.\n>\n>We also believe that extending our method to other similar tasks and datasets can help solve overwhelming tasks such as reassembling shredded documents or restoring fragmented objects and images, and potentially even more exciting future directions, such as designing electrical boards based on the provided components.\n---\n**2. Non-learning based methods for Full MagicPlan (RPLAN).**\n\n>As highlighted by Shabani et al. [20], their method needs several minutes to process a small-scale house. However, their processing time explodes exponentially to approximately one hour for a house with seven rooms and even a few days for a house with ten rooms.\n>\n>Analytically, suppose each room has two doors, the number of possible room arrangements by aligning pairs of doors is equal to the number of spanning trees in a complete graph with $n$ vertices, which is $n^{(n-2)}$.  The count becomes approximately 1k, 262k, and 100 million for a floorplan of 6, 8, and 10 rooms, respectively. This demonstrates the severe scalability challenge of non-learning based methods.\n\n**3. The canonical diffusion denoising probabilistic model requires quite some denoising steps.**\n\n>We thank for the great comment. Our current method utilizes the standard diffusion denoising probabilistic models, which are slow and require many sampling steps. However, we rather consider the use of diffusion models to be a distinct advantage of our approach. This is because any advancements in accelerated sampling techniques and efficient diffusion models, which are currently a vibrant area of research, will be directly applicable to our method. This aspect is beyond the scope of our paper due to space constraints and the diverse array of available sampling approaches.\n\n---\n\n**4. It seems better to derive more analysis on ablation studies.** \n\n>We thank for the interesting questions. We will add the following analysis to the main paper. \n>- Matching loss on doors and corners: One reason can be that the corners of rooms within the floorplans do not necessarily align. For instance, a wall from the dining hall might divide two rooms, thereby not conveying significant high-level information. In contrast, doors hold more informative value as they connect the rooms.\n>- Feature redundancy: While the redundancy in the diffusion process can be considered as adding more randomness and possibly more augmentation, as mentioned by the reviewer, it also serves us in two ways. 1) It helps in the model's architecture design by providing the ability to establish explicit connections among corners within the transformer, for example by utilizing Polygon Self Attention. 2) During the final prediction stage, this redundancy can be beneficial by allowing the averaging or voting of predictions akin to an ensemble model.\n\n___\n\n**5. Comparison with Lambert et al. [21] for Layout arrangement.**\n\n>While the high-level problem is the same, Lambert et al. [21] assume that the panorama images are given and have some overlaps. They leverage this information to generate Aligned Bird's Eye View Texture Maps, subsequently employing these maps for global pose estimation. However, this is not possible in our targetted application domains, because in a mobile application such as MagicPlan, panorama images cannot be uploaded to the cloud due to the privacy/security concerns as well as to minimize the amount of data transfer from a mobile device (see L296-300). Our approach requires only room layouts that are extremely compact and does not convey any private information.\n\n---\n**6. How does the method cope with the variability of the DDPM output during the evaluation?**\n\n>We thank for the question and will clarify the following in the paper. We run our system five times and report the mean.\n\n---\n**7. Typo in Tabel. 2 Caption.** \n>We thank for the typo check and will make a correction."}]

Author rebuttal:
We thank all the reviewers for their valuable and insightful comments. We are also grateful to the reviewers for their positive comments on our work. We have addressed the reviewers points in our individual responses to each reviewer, and please let us know if there are any new questions.


Meta-review: The paper received unanimous positive feedbacks. The paper is well-written and the problem is novel. Some issues related to clarity and applications are resolved after the rebuttal.
Decision: Accept (spotlight)