Leveraging Pretrained Diffusion Models for Zero-Shot Part Assembly

1 Introduction

3D part assembly autonomously assembles unordered 3D pieces into a realistic, complete object by predicting the rotations and translations of each piece. This research topic has drawn great attention in the field of robots in recent years, as it plays a crucial role in advancing robotic manipulation and automation [Chervinskii et al., 2023; Ghasemipour et al., 2022; Zhan et al., 2020; Zhang et al., 2022; Gao et al., 2024].

3D part assembly is challenging because of the intricate geometries and various possible assembly combinations. The existing approach to 3D part assembly relies on training machine learning models with extensive manually annotated data, including rotations and scalings. However, the high cost of data collection makes it impractical to create datasets for each task, limiting supervised methods to well-resourced domains like common datasets. This question drove us to search for new methods to reduce reliance on manual labeling.

Diffusion models are a recent class of likelihood-based generative models that model data distributions through an iterative noising and denoising process [Ho et al., 2020]. Following this, diffusion-based distillation models [Wang et al., 2024b] have demonstrated significant high-fidelity 3D content generation capabilities, highlighting both their theoretical robustness and practical applicability in generating a large amount of complex 3D contents. Meanwhile, some studies have also demonstrated that a pre-trained diffusion model, leveraging its density estimates, can be transferred to handle various zero-shot tasks, including classification [Li et al., 2023; Zhang et al., 2024a], vision-language modeling [Liu et al., 2025; Liu et al., 2024; Liu et al., 2023b], and singing voice synthesis [Zhang et al., 2024d; Guo et al., 2025; Zhang et al., 2024c; Zhang et al., 2024e]. Density estimation refers to the distribution of particles in space as they evolve over time during a diffusion process. These works further inspire us to explore how to distill the necessary pose transformations in assembly tasks using existing diffusion models.

In this paper, we propose a new algorithm for aligning density estimates to pose transformations. Specifically, we first introduce noise to a shape that has not been correctly positioned. This perturbed shape is then input into the diffusion model. The objective at this stage is to transform the disordered components into a distribution that is suitable for the diffusion model. By utilizing the denoising process, we can obtain a new point cloud that is closer to an accurate chair shape. However, it is important to note that this new point cloud does not represent a rigid transformation compared to the previous point cloud. To address this issue, we employ the ICP algorithm to align each part as closely as possible. By iteratively repeating this process, we can utilize the diffusion model to convert disordered parts into a complete shape, thereby accomplishing the entire assembly process. Since this process is performed explicitly, it allows us to apply direct pull-in or push-away operations for overlapping or distant parts, which is nearly impossible to achieve with other methods. To validate our method, we employed four network architectures to predict rotational and translational transformations of parts. These baselines rely on the Shape Chamfer Distance (SCD) for supervised learning, aiming to approximate ground truth derived from reference samples generated by the diffusion model. Quantitative and qualitative results indicate that our method not only outperforms all baseline approaches in zero-shot settings but also surpasses some supervised techniques, underscoring its potential for practical applications.

The contributions of our paper can be summarized as follows:

• We propose the first zero-shot assembly method that utilizes density estimates from a diffusion model to achieve continuous and smooth transformations of parts, thereby coherently assembling multiple parts. Theoretical analysis within the paper supports the efficacy of this approach.
• We additionally introduce a push-away strategy to mitigate collisions between parts.
• Results show that our method outperforms all baselines in zero-shot settings and even some supervised approaches, highlighting its practical potential.

2 RelatedWorks

2.1 3DAssemblyModeling

Estimating object pose has been a key research focus for decades. In the early research, Yoon et al. [2003] used visual sensors and neural networks for robotic assembly. Later, graph models were employed to capture semantic and geometric relationships among shape components, enabling advancements in assembly-based shape modeling [Zhan et al., 2020; Jaiswal et al., 2016], while a progressive strategy leveraging the recurrent graph learning framework was explored in [Narayan et al., 2022]. To explore the diversity of assembly outcomes, several authors propose treating parts' poses as a distribution and achieving part assembly through a diffusion process involving noising and denoising [Xu et al., 2024; Cheng et al., 2023]. Furthermore, innovations in network architecture have been advancing concurrently. For instance, Zhang et al. [2024b; 2025] leverage the Transformer framework [Vaswani, 2017] to model structural relationships. Building on this, Gao et al. [2024] introduces hierarchical assembly to tackle the challenges associated with managing numerous parts. Unlike the aforementioned works that rely on manual annotations of each part's rotation and translation, our study aims to explore a novel approach to extracting the necessary pose transformations for assembly tasks. Specifically, we investigate how existing diffusion models can be leveraged to achieve this goal, thereby reducing the dependency on labour-intensive manual labelling.

2.2 DiffusionModel

Diffusion models operate in two steps: adding noise to destroy data structure and reversing this noise to reconstruct it. This enables them to model target distributions and generate diverse content, including images [Saharia et al., 2022], videos [Wang et al., 2024a], 3D objects [Peebles and Xie, 2023], and audio [Liu et al., 2023a]. Recent studies suggest that diffusion models encode semantic and grouping information, leading to two main research directions. The first research direction leverages the internal representations of diffusion models for various discriminative tasks, requiring minimal additional training. These tasks include zero-shot classification [Li et al., 2023], label-efficient segmentation [Baranchuk et al., 2021], and open-vocabulary segmentation [Karazija et al., 2023]. The second research direction focuses on generative tasks, such as bridging 2D diffusion models and 3D generation through Score Distillation Sampling (SDS). Methods like DreamFusion [Poole et al., 2022] align 3D representations with text prompts, while later works enhance visual fidelity using strategies like coarse-to-fine optimization [Lin et al., 2023] and multi-view consistency [Hu et al., 2024]. These advances highlight the versatility of diffusion models in blending discriminative and generative capabilities. Our work builds upon the generative approach, introducing a theoretically sound and interpretable method to tackle the zero-shot assembly problem effectively.

3 Methodology

In this section, we will first provide a formal symbolic definition of diffusion models (Sec. 3.1). Next, we will introduce the zero-shot method proposed in this paper (Sec. 3.2). Finally, based on our method, we will present a new approach to mitigate part overlap (Sec. 3.3).

3.1 DiffusionModel Preliminaries

The diffusion model [Song and Ermon, 2019; Luo and Hu, 2021] is a likelihood-based generative model, designed to learn the data distributions. Starting from an underlying data distribution q(x), the model applies a forward process that progressively adds noise to a data sample x, creating a sequence of latent variables {xz}Zz=1 governed by Gaussian transition kernels q(xz|xz-1). At each time step z, the marginal distribution of xz is defined as:

where σ2 z + α2 z = 1, with σz gradually increasing from 0 to 1. This ensures that q(xz) converges to a Gaussian prior distribution N(0, I) as z approaches Z. Thus, q(xz) converges to a Gaussian prior distribution N(0, I).

The reverse process, which corresponds to the generative process, is designed to reconstruct the original data from a sequence of noisy observations. The conditional distribution pϕ(xz-1|xz) at each time step z is modeled as a Gaussian with mean µϕ(xz, z) and covariance variance Σϕ(xz, z):

To ensure that the model can accurately reconstruct the original signal as it approaches the end of the generation process, Σϕ(xz, z) is typically designed to decrease as z decreases. This reflects the intuition that the model's confidence in predicting the next state should increase as it gets closer to the original data point.

Specifically, Σϕ(xz, z) can be parameterized or fixed according to a schedule that depends on the time step z. In practice, this variance term may be simplified to depend only on z, for instance, by setting it proportional to the pre-defined noise scale σ2 z :

where σ2 z is part of a predefined noise schedule that increases over time during the forward diffusion process and consequently decreases during the reverse generative process. A linear noise schedule could be defined as:

As such, when z is small, indicating that we are close to the final generation step, σ2 z is also small, leading to a smaller Σϕ(xz, z). This design choice ensures that the model exhibits higher certainty in its predictions as it nears the reconstruction of the original data, thereby enhancing the stability and quality of the generated samples.

3.2 Diffusion Based Iterative Zero-Shot Assembler

Denote the input point clouds as P = {Pi | i = 1, . . . ,N}, where Pi ∈ Rd×3 corresponds to the i-th part of the 3D shape, consisting of d points in the 3D space. In zero-shot task, each part of the point cloud Pi has a corresponding rigid transformation, described by a quaternion quati ∈ R4 and a translation vector transi ∈ R3 , which represent the rotation and translation of the part. The goal of this task is to predict the pose parameters (quaternion quat and translation vector trans) of the test samples without pose information during training.

To match the current shape to the diffusion models's requirements, we introduce Gaussian noise to the current shape:

where t is the iterative step of our method, z is the time step in diffusion model ϵ.

By utilizing the denoising process, we can obtain a new point cloud that is closer to the shape's distribution:

where c corresponds to the prompt label associated with the input sample. Subsequently, to satisfy the requirements of rigid transformations, we employ ICP to obtain the vector of rotation and translation. We then apply the transformations to the input point cloud Pt to obtain the updated poses, which are then utilized as the input for the next iteration.

The theoretical justification for using ICP is detailed in Section 4. By iterating the above process, we can utilize the diffusion model ϵ to convert disordered parts into a complete shape, thereby accomplishing the entire assembly process.

3.3 Collision detection and handling

Given the explicit nature of our method, it facilitates the direct application of pull-in or push-away operations for either overlapping or distant parts. This strategy is very difficult to implement in existing methods due to their poses being implicitly generated by the model. To describe the pushing behavior of Pi in a point cloud P to reduce overlap with Pj (i ̸= j), the overlap is quantified using C(Pi,Pj), which counts coincident points. The indicator function I(C(Pi, Pj) < threshold) determines whether the overlap is below a predefined threshold. The centroids of Pi and their intersection region are denoted as ¯ CPi and ¯ Cintersect, respectively. The displacement required to separate Pi is given by:

Here, s specifies the sign of the movement. This approach computes the necessary displacement direction and distance to reduce the overlap between Pi and Pj .

4 The Theory of Zero-Shot Assembly

As mentioned previously, part assembly seeks to optimize the rotation qi and translation ti of each part Pi to transform the unordered input into a coherent realistic object. Let the qi and ti be represented by an optimizable transformation matrix Ai, then the assembly process can be formulated as:

Figure 2: Visual comparisons demonstrating our superior assembly performance over baseline methods on PartNet. The first column

shows our input at the Excessive level, while the last column presents reference samples obtained through diffusion sampling.

where A = {Ai | i = 1, ...,N}, g(∗) denotes the matrix multiplication with the unordered point cloud. For previously supervised part assemblers, the optimization of A is quite straightforward:

where E denotes a set of distance functions and Pgt is the ground truth. However, optimizing A is non-trivial in our case, where no supervised data is available. Therefore, instead of forcing g(A) to fit a determined object, we tend to make the generation of g(A) looks like a realistic object, i.e. a sample from the distribution of the real object. Inspired by Poole et al. [2022], we leverage a pre-trained diffusion model for 3D point cloud generation, which implicitly captures the distribution of point clouds in real-world objects. Then we optimize over A so that g(A) looks like a sample from this frozen diffusion model. This is achieved through a Score Distillation Sampling (SDS) loss [Poole et al., 2022]:

As shown in Fig.1:

Substituting Eq. 10 and Eq. 11 into Eq. 9, we get:

In the equation above (Eq. 12), since g(∗) represents matrix multiplication, ∂Pout ∂A corresponds to the coordinates of the points in Pt, which are constant three-dimensional vectors; w(z) is a constant scalar in practice, which will be explained in the following section. Therefore, the remaining term Pt -P ∗ governs the descent process of the LSDS. If we can accurately estimate the transformation fromPt to P ∗ , the optimization process will converge directly.

In practice, we utilize the ICP algorithm to estimate the transformation between Pt and P ∗ , as shown in Fig. 1. It is worth noting that the smaller the change in the shape of each part in Pt and P ∗ , the more accurate the transformation obtained by the ICP algorithm. We minimize the shape variation between Pt and P ∗ by controlling the magnitude of noise added and removed during the forward and generation processes. Recall that the step size determines the noise (Eq. 1 and 3): (1) In the forward process, smaller step sizes reduce the Gaussian noise variance, bringing Pt closer to Pt,z; (2) In the generation process, smaller step sizes reduce denoising variance, making Pt,z closer to P ∗ . Therefore, we fixed the time step z to a small value, typically 2 or 4, to obtain more accurate ICP estimates. Fig. 3 shows the assembly result under different z with the same iterations, a smaller z significantly improves the realism of the results. We finally apply the transformations obtained from the ICP algorithm to Pt to generate the result of this iteration Pt-1 and use it as the input for the next iteration. With each iteration, the diffusion model helps bring our results closer to real-world objects.

5 Experiments

5.1 Datasets, Baselines, andMetrics

Dataset. We conduct experiments on the Chair subset of PartNet [2019], a large-scale dataset with fine-grained part-level annotations. We follow the official train/val/test splits: the training set is used to train a diffusion model for 3D point cloud generation, and the test set is used for zero-shot assembly. Part counts range from 2 to 20. To evaluate robustness, we introduce four noise levels: slight, moderate, substantial, and excessive (see Fig. 6). The diffusion model is adopted from [Luo and Hu, 2021].

Baselines. We compared our approach with Complement [2017], DGL [2020], IET [2022], HPA [2024], and Simple. Among them, Simple is the Baseline we designed, which utilizes seven trainable parameters to represent the rotational and translational transformations of parts. Simple usually outperforms other baseline methods in practice.

Metrics. We evaluate using Part Accuracy (PA) and Shape Chamfer Distance (SCD) from Zhan et al.[2020], as well as Rotation and Translation RMSE (RMSE(R), RMSE(T)) from Sell´an et al.[2022]. PA assesses per-part precision, SCD measures overall shape quality, and RMSE(R/T) quantify rotation and translation accuracy.

Fair Part Accuracy (fPA). Vanilla PA is determined by the Chamfer Distance between components with identical tensor indices. This metric serves as a criterion for evaluating assembly precision. However, certain components, such as stool legs, are permitted to be positioned in regions with inconsistent indices, as illustrated by GT and Ours in Figure. Therefore, we propose the concept of Fair Part Accuracy (fPA). Given two point clouds Ppred = {p[i], i ∈ {1, 2, . . . , N}} and Pgt = {g[i], i ∈ {1, 2, . . . , N}}, CD represents Chamfer Distance calculation. we formally define:

Next, we use Pfair gt to replace Pgt.

We define the accuracy as:

where thre = 0.01, which is a parameter inherited from previous work [2020]. 1 denotes an indicator function that equals 1 if the condition inside is met, and 0 otherwise.

5.2 Experiments Results and Analysis

As demonstrated in the first four rows of Table 1, assembly performance declines with increasing noise intensity, thereby validating our noise level designations. The slight noise level evaluates the ability of our method to converge close to the ground truth. Conversely, the excessive noise level, characterized by randomly dispersed point clouds, tests the extremes of performance for both baselines and our method. The intermediate moderate and substantial levels further substantiate the efficacy and rationality of our noise addition strategy.

We evaluated our method against various baselines, as shown in Table 1, Fig. 2, and the Appendix. Our approach outperforms current methods in addressing the zeroshot challenge. All of our metrics outperform existing methods, except for SCD. This is expected since our method emphasizes density estimates from a diffusion model rather than sampling complete shapes. Therefore, our assembled outputs conceptually resemble chairs instead of precisely replicating chair-shaped point clouds. This distinction is illustrated in Fig. 4, which visualizes samples generated by Simple under two different random seeds. Furthermore, compared to IET, Simple performs worse on the SCD, highlighting the advantage of Transformer models in encoding complex structures. Additionally, we tested a carefully designed model, HPA, whose performance is significantly impaired when trained exclusively with the SCD. Without any prior information on part

e

Table 1: Quantitative evaluation on zero-shot scenario. Underline/bold fonts highlight the suboptimal/best approach. Our approach

outperforms current methods in addressing the zero-shot challenge.

poses, all baselines demonstrate notably poor performance, as further analysis of their training loss functions reveals why they fail in zero-shot scenarios (details in the Appendix). In the Appendix, we also present an experiment designed to illustrate both the effectiveness of the proposed method and its limitations on challenging samples. To evaluate the proposed method under varying levels of task complexity, we conducted experiments with different numbers of components in the Appendix. As well as more visual results and the details of our experiments.

5.3 Comparisons with supervised scenario

As indicated in Table 2, our work achieves comparable results by the early supervised learning method: Complement. This finding underscores that our method can deliver competitive outcomes even compared to pose-accessible supervised learning. While our work may not yet achieve the performance of existing well-designed supervised learning methods, we hope it offers insights that may contribute to future research in zero-shot learning.

5.4 Ablation Study

In this work, we propose an improved approach to address the collision issue that arises when identical parts are placed in the same position. To mitigate this issue, we introduce an explicit pushing-apart operation. Traditional model-based training methods struggle to achieve this directly, as the model generates the predicted poses of parts, and adjustments can only be made by tuning the model parameters, which limits operational flexibility. However, in this study, we innovatively incorporate an explicit pushing-apart operation into the fragmented images without simplifying the problem, achieving near-perfect reconstruction for 3×3 image puzzles.

original method. This operation effectively separates parts, thereby reducing collision issues. The experimental results, as shown in Fig. 5, demonstrate the effectiveness of the proposed method.

5.5 Tranfer our work to 2DDiffusionModel

As analyzed in Section 4, our method is not limited to specific types of diffusion models. A pretrained 2D image diffusion model, in theory, can also evaluate the quality of assembly results by differentiably rendering 3D components into 2D space and feeding them into the model. However, in practice, this approach is challenging due to the difficulty of propagating 2D signals back to 3D point clouds and pose parameters. To test this potential, we conducted a simplified experiment where all parts shared a common rotational perturbation. This perturbation was optimized using Eq. 9. We utilized stable diffusion 2.1 [Rombach et al., 2022] with the prompt "a picture of colorful chair" to generate a realistic chair shape. As shown in Fig. 8, we demonstrate the visualization of the assembly process. These images are processed through microrendering, serving as inputs for the diffusion model to obtain SDS loss. This experiment demonstrates that 2D diffusion models have inherent assembly potential, though sophisticated methods are required to fully utilize it. More details in the Appendix.

5.6 Transfer our work to 2D Image Reassembly

In addition to exploring the potential of 2D diffusion models for 3D part assembly tasks, we also examined their capability in 2D image reassembly. Our model can almost perfectly reconstruct 3×3 image fragments without oversimplifying the challenge. 2D Image Reassembly involves reassembling cropped segments of a 2D image [Scarpellini et al., 2024]. In this experiment, no special architecture was designed; instead, we implemented a classifier utilizing two CNN layers. An MLP was implemented to predict the correct position of each sub-image within the whole picture. Training followed the method outlined in Eq. 9, using stable diffusion 2.1 [Rombach et al., 2022] with the prompt "a picture of chair". As illustrated in Fig. 7, our approach successfully handles 2D image reassembly challenges up to a complexity of 4×4. More details in the Appendix.

6 Conclusion

In this work, we introduced a novel zero-shot assembly method that leverages the inherent assembly capabilities of general-purpose diffusion models to generate continuous rigid transformations for object assembly without prior training on specific shapes or configurations. It uncovers the implicit assembly abilities of general models, enabling assembly tasks even with previously unseen data. Our work is the first zero-shot part assembly framework, which aims to harness existing extensive work on diffusion models to achieve assembly at virtually no additional cost, which represents a meaningful and valuable contribution.

7 Limitation and FutureWork

The bottom row of Fig. 2 exemplifies a failure case, revealing challenges in accurately placing overlapping parts. Especially when they are far from their GT positions, using the push operation cannot accurately place overlapping parts. Adjusting random inputs can mitigate this issue, but it is not robust enough. In the future, we plan to investigate an interpretable approach to reposition misplaced parts. This requires us to better explore the underutilized assembly knowledge inherent in general models and to be able to identify which positions have vacancies, thus allowing for the effective transfer of parts. Our goal is to improve the model's performance in assembly tasks without extensive additional costs or extensive supervised learning.

Acknowledgments

This work was supported by the Zhejiang Provincial Key Research and Development Project (2023C01043), Engineering Research Center of Integration and Application of Digital Learning Technology, Ministry of Education, and the Academy of Social Governance, Zhejiang University. Ruiyuan Zhang and Qi Wang contributed equally to this work. Please ask Prof. Wu (chao.wu@zju.edu.cn) for correspondence.