SE(3)-Equivariant Diffusion Models for 3D Object Analysis

1 Introduction

3D vision is a crucial branch of computer vision that focuses on the comprehension and analysis of 3D objects, scenes, and motion. Humans instinctively solve complex spatial problems by leveraging spatial relationships and transformation equivariance to understand how various parts fit together. This capability is particularly valuable in fields such as archaeology, medicine, and biomedicine, where reconstructing fragmented 3D objects requires both accurate identity recognition and precise pose estimation. However, predicting accurately pose is a challenge for conventional invariant networks [Xie et al., 2024], which produce constant outputs regardless of input transformations, lack the flexibility to capture spatial variability. This limitation has driven the development of advanced networks, such as equivariant networks, which can adapt to diverse transformations.

∗Corresponding author

Equivariant networks ensure that a transformation applied to the input results in a correspondingly transformed output. This property facilitates the learning of part transformation and improves generalization to unseen data with various spatial configurations. Existing methods for achieving equivariance in 3D vision can be broadly categorized into equivariant kernel networks and equivariant tensor networks. Equivariant kernel networks [Hoogeboom et al., 2022], [Hutchinson et al., 2021] extend convolution kernels using higher-order representations, such as spherical harmonics or Lie group representations [Yu and Sun, 2024]. While equivariant tensor networks [Son et al., 2024], [Assaad et al., 2023], [Lei et al., 2023] enrich scalar features by evolving them into tensorvalued representations, enhancing the understanding of input.

Despite these advances, predicting part poses of fragmented 3D objects remains challenging. Single-layer equivariant networks struggle to simultaneously capture fine-grained part-level pose information and assemble these parts into complete objects. To address this limitation, hierarchical structures, such as equivariant graph networks [Scarpellini et al., 2024], have been proposed. These networks [Satorras et al., 2021], [Le et al., 2022], [Kofinas et al., 2024], [Du et al., 2022], [Meng et al., 2024] use message-passing mechanisms to maintain spatial relationships and capture geometric relationships between parts. However, they often fall sort when parts undergo different transformations, such as rotations.

To overcome this challenge, we draw inspiration from diffusion models, which excel in handling uncertainty and iterative refinement. Diffusion models, successful in tasks such as molecule generation [Guan et al., 2023], robotic manipulation [Ryu et al., 2024] and trajectory prediction [Chen et al., 2023], [Liu et al., 2024], perform iterative denoising to recover information from noisy data. Unlike NeRF [Zhou et al., 2025] or Gaussian splatting [Ma et al., 2024], which focus on fixed viewpoints or scenes, diffusion models are wellsuited to tasks requiring iterative refinement, such as 3D part pose learning. By treating random transformations as noise, we conceptualize the learning of transformed 3D part poses as a denoising process [Bansal et al., 2024], where the goal is to recover original poses and assemble parts into objects.

Building on this intuition, we propose a novel framework that integrates diffusion models with SE(3)-equivariance, SE(3)-equivariance maintains consistency under rotations and translations in 3D Euclidean space. Equivariance ensures consistent transformations across parts, while diffusion models provide the flexibility to generate and refine poses in diverse spatial configurations. However, a key challenge is the fixed input size of conventional diffusion models, which limits their adaptability to varying part counts. Inspired by advancements in elastic diffusion for 2D images [Zheng et al., 2024], [Haji-Ali et al., 2024], we extend these principles to the 3D domain, enabling our model to process inputs with varying part quantities while maintaining equivariance.

We evaluate our approach on the public 3D part dataset Breaking Bad [Sell´an et al., 2022], a real-world archaeological dataset RePAIR [Tsesmelis et al., 2024] and a self-built dataset of fragmented 3D mannequin. These datasets allow us to train geometry-based networks without relying on semantic labels, providing the robustness of our framework. In summary, we make the following contributions:

• We propose a novel Lie algebra mapping to linearize the processes of noise addition and removal, enabling the transformation of 3D input data into a representation compatible with diffusion models.
• Building on this, we design an elastic diffusion model capable of handling arbitrary numbers of 3D parts, effectively adapting to varying input sizes while accurately denoising and predicting the pose states.
• Comprehensive experiments conducted on the Breaking Bad dataset, real-world RePAIR dataset and a selfconstructed 3D mannequin dataset demonstrate that our method significantly outperforms state-of-the-art approaches in terms of root mean square error, mean absolute error, and part accuracy for 3D part assembly tasks.

2 RelatedWorks

2.1 Equivariant Networks

Our work focuses on SE(3)-equivariant networks, formally, an SE(3)-equivariant function f satisfies g(f(X)) = f(g(X)), for all g ∈ SE(3) applied to an input X. Group equivariant convolution extends traditional convolution to group convolution under discrete symmetric group actions. For instance, Thomas et al. [Thomas, 2019] and Fuchs et al. [Fuchs et al., 2020] use spherical harmonics as high-order filters to design equivariant networks. However, these methods often lack interpretability and flexibility due to architectural constraints [Son et al., 2024].

To overcome this limitation, recent works [Deng et al., 2021], [Assaad et al., 2023],[Lei et al., 2023] extend scalar features to tensor-valued representations, enabling network to model complex transformations. Son et al.[Son et al., 2024]

further map 3D point coordinates into a high-dimensional sinusoidal feature spaces for enhanced shape compression.

However, these single-layer equivariant networks struggle with tasks requiring the integration of individual part poses into a cohesive structure. To address this, researchers are exploring hierarchical structures for equivariant networks that can better model these complex relationships.

2.2 Hierarchical Networks and DiffusionModels

Hierarchical graph structures have been proposed to model spatial relationships among parts. Graph neural networks, which are typically permutation-equivariant [Huang et al., 2024], maintain topological relationships through messagepassing mechanisms. For example, Kofinas et al. [Kofinas et al., 2024] represent neural networks as hierarchical computational graphs, preserving equivariance and spatial structure.

Diffusion models have also been explored for hierarchical tasks. Gianluca et al. [Scarpellini et al., 2024] propose a graph-based diffusion model for refining 3D part poses, treating fragments as nodes in a spatial graph. Similarly, Wang et al. [Wang et al., 2024] extend the PuzzleFusion [Hossieni et al., 2024] concept to 3D, iteratively refining 6-DoF alignment parameters. These models demonstrate the potential of diffusion-based approaches for iterative pose refinement.

However, conventional diffusion models are typically trained on fixed-size inputs, they struggle with varying aspect ratios during inference. Therefore, designing a flexible diffusion model to handle diverse input size is a new challenge.

2.3 Elastic Feature Selection

Recent advancements in 2D diffusion models address the challenge of variable input sizes [Zheng et al., 2024], [Podell et al., 2023]. For instance, ElasticDiffusion [Haji-Ali et al., 2024] decouples global and local content generation, enabling robust synthesis across resolutions. Inspired by these methods, we design elastic diffusion models for 3D object analysis, allowing the network to adapt to varying part counts while preserving equivariance. This flexibility is crucial for assembly tasks, where input configurations can vary widely.

3 Equivariant DiffusionModels

This section introduces the proposed equivariant diffusion framework, which aims to denoise pose transformations and reassemble fragmented 3D parts into a coherent object. The framework consists of three main components: (1) an equivariant feature representation module to extract translational and rotational equivariant features, (2) a Lie algebra mapping to linearize the addition and removal of noise in the transformation matrices, (3) an elastic diffusion model designed to handle diverse input sizes while ensuring accurate denoising.

3.1 Equivariant Feature Representation

To obtain equivariant features for fragmented parts, we first translate the parts {P k }Kk=1 such that their centers lie at the origin, ensuring that the input parts { ¯ P k }Kk=1 are translationally equivariant. Each part P k is a sampled point cloud containing 1000 points.

Then a rotation-equivariant encoder is employed to extract features and predict pose matrices {M k 0 }Kk=1 . In this paper, we utilize the encoder proposed in [Xie et al., 2024], which leverages an equivariant masked autoencoder to learn robust features for 3D objects. This encoder's feature learning capability has been validated through its reconstruction performance on various datasets.

The predicted pose matrices are then input into the diffusion model. To disrupt the input, we replace Gaussian noise with transformed matrices during the forward process. During the reverse process, the diffusion model denoises the matrices to estimate the original transformations, enabling the reassembly of fragmented parts into a complete 3D object.

3.2 Lie AlgebraMapping

Directly replacing Gaussian noise with transformation matrices Φknoise ∈ R3×3 often leads to degeneration in the object's dimensions, as shown in the first row (orange object) of Figure 1. This degeneration occurs because traditional diffusion models rely on linear noise addition/removal, which is incompatible with the nonlinear nature of transformation matrices.

To address this issue, we map the transformation matrices to the Lie algebra space so(3) [Batatia et al., 2023], where nonlinear operations can be transformed into linear ones (as shown in the sencond row, the blue object, of Figure 1). Specifically, we use the logarithmic map log : SO(3) → so(3) to project rotation matrices Φknoise into vectors φknoise :

where u = [ux, uy, uz]⊤ is a unit vector representing the rotation axis, θ ∈ [0, π] is the rotation angle. φknoise is a 3 × 3 skew-symmetric matrix, (·)∨ is the mapping from a skewsymmetric matrix in so(3) to a vector. The inverse mapping is performed via the exponential map exp : so(3) → SO(3):

By leveraging the Lie algebra mapping to linearize the noise addition and removal processes, the part pose features are transformed into a representation compatible with diffusion models. Building on this foundation, we propose an elastic diffusion model designed to handle diverse input sizes and perform effective denoising and pose prediction.

3.3 Elastic DiffusionModels

Before feeding these mapped vectors into diffusion model, we ensure compatibility with varying input sizes by applying padding and cropping operations. The padded vectors are:

where φ = {φGT , φnoise}, Z is a zero matrix for padding. Padding function fP ensures that the number of ˆ φ is always N, cropping function fC ensures that the output size matches the original part count K. This padding-cropping strategy guarantees consistent input sizes N while retaining the focus on the actual parts K.

The process of the diffusion model for 3D part pose learning is illustrated in Figure 2. Each row corresponds to the transformation and recovery of a single fragmented part. The framework consists of a forward diffusion process and a reverse diffusion process, which ensure accurate reconstruction of the original pose for each part, enabling effective 3D part reassembly.

Forward process. In Lie algebra space, noise is added linearly to the transformed pose vectors:

where ˆmkt represents the noisy pose at step t, M kGT =M k 0 is the original pose state.

Reverse process. The model estimates and removes noise to recover the original pose:

where ˆ ϕk,t noise is the estimated noise output by our models that has to be removed from ˆmkt at timestep t to recover ˆmkt-1 . The denoised Lie algebra element is then mapped back to SO(3) space:

Finally, the denoised pose matrix ˙M k 0 = { ˙Rk 0 , ˙T k 0 } should correspond to the original pose matrix, ensuring accurate reassembly of the fragmented parts.

Loss function. To optimize the diffusion models, we introduce three loss functions: translation loss Ltrans, rotation loss Lrot, and point loss Lpoint. The total loss is:

where {λi}3i=1 are weights. Each term ensures precise reassemble of translation, rotation.

Translation loss. The translation loss computes the distance between the ground truth T0 and the predicted ones ˙T0:

where T0 = {T k 0 }Kk=1 is the ground truth of translation, ‖·‖22 is the L2 loss.

Rotation loss. The rotation loss measures the geodesic distance between ˙R0 and R0:

Point loss. We further use Chamfer Distance fcd(·) to jointly measure the difference by supervising the reassembled pose of point cloud:

4 Experiments

4.1 Dataset and Implementation Details

Datasets. This study focuses on denoising poses to recover the original part states and reassembling transformed parts into a coherent object. To evaluate the proposed approach, we conduct experiments on the publicly available Breaking Bad dataset [Sell´an et al., 2022], a real-world RePAIR dataset [Tsesmelis et al., 2024] and construct a self-built 3D mannequin dataset to assess the robustness of our method.

The Breaking Bad dataset [Sell´an et al., 2022] consists of approximately 10,000 meshes derived from PartNet [Mo et al., 2019] and Thingi10k [Zhou et al., 2016]. Specifically, we select samples from the everyday subset and 6 artifact categories (sculpture, spiral bulbs, rabbit, frog, sofa and boy) for experiments. The everyday subset contains 20 object categories, each containing fragments with varying quantities. The RePAIR dataset [Tsesmelis et al., 2024] consists of over 1,000 verified fragmented parts, it contains detailed 3D scans of the fragments, the fragments and fractures are realistic, caused by a collapse of a fresco during a World War II bombing at the Pompeii archaeological park. While the Breaking Bad and RePAIR datasets provides diverse object categories, it lacks sufficient variation in human-like object structures, which are essential for evaluating real-world applications such as pose estimation and archaeological human fragment restoration.

To address this limitation, we construct the 3D mannequin dataset, focusing on human-like structures with predefined semantic parts. This dataset consists of 133 simulated mannequin samples collected from ShapeNet [Yi et al., 2016]. These samples are preprocessed by SubdivNet [Hu et al., 2022] to ensure smoothness and uniformity, resulting in mesh with 16,384 faces each. The mannequins are segmented into distinct fragments based on semantic labels, with each fragment stored in OBJ format. We also create five subsets with varying fragment proportions. For example, a sample in the Mannequin 4Parts subset includes a head, torso, arms, and legs, as illustrated in Figure 3.

Implementation details. All experiments are performed on a Linux workstation equipped with an NVIDIA RTX 4090 GPU. We use Noise Conditional Score Networks (NCSNs) [Song and Ermon, 2019], [Croitoru et al., 2023] as the backbone of the diffusion models. For the 3D object reassembly task, we adopt the encoder proposed in [Xie et al., 2024] to extract features from each part and predict the corresponding poses. These predicted poses are then fed into the diffusion models for denoising, allowing the refinement of poses toward their original states.

Optimization is performed using the Adam optimizer with an initial learning rate of 1e-3 and a cosine learning schedule. The batch size is set to 16, and the total number of iterations is 2000, with a consistent input size N = 30. To further improve performance, we incorporate a Chamfer Distance loss term, as proposed in [Sell´an et al., 2022], which enhances the model's capability to minimize geometric discrepancies.

Evaluation metrics. We evaluate our method using multiple metrics, including root mean square error (RMSE) metrics [Sell´an et al., 2022], RMSE(R) for rotation and RMSE(T) for translation, and the mean absolute error (MAE) for both rotation and translation. Additionally, part accuracy (PA) [Sell´an et al., 2022] employed to assess the precision of parts.

4.2 3DObject Reassembly Task

The quantitative evaluation of our method, compared to stateof-the-art methods, on the 3D Mannequin nParts datasets is presented in Table 1. The results indicate that while the performance of our method decreases with an increasing number of parts, it consistently ourperforms SE(3)-Equiv [Wu et al., 2023], Jigsaw [Lu et al., 2024] and DiffAssemble [Scarpellini et al., 2024]. Visualization results are shown in Figure 4, and the iterative refinement process is illustrated in Figure 6.

Qualitative results demonstrate that our method effectively denoises the pose noise and iteratively predicts the original poses. Notably, we observe that SE(3)-Equiv and Jigsaw struggle with larger parts, while DiffAssemble tends to shift all parts toward the center, as shown in Figure 4. It is worth noting that the runtime of the method proposed in this article is measured from the moment the obtained poses are fed into the diffusion model. This is primarily done to evaluate the diffusion model's ability to learn poses.

We also evaluate our method on the Breaking Bad and RePAIR datasets, we select all categories from the everyday subset and 6 categories from artifact subset of Breaking Bad dataset as experimental subjects. Quantitative results are summarized in Table 2 and 3.

These results highlight superior performance of our method across various metrics, emphasizing its accuracy and robustness in 3D object reassembly task. Although our method requires more time compared to SE(3)-Equiv [Wu et al., 2023], CCS [Zhang et al., 2024] and Jigsaw [Lu et al., 2024], its superior performance justifies the additional computational cost. Visualization results are shown in Figure 6, and the iterative process is illustrated in Figure 7.

4.3 Ablation Studies

To evaluate the contribution of different modules, we conduct ablation studies on Mannequin 6Parts, Breaking Bad and RePAIR datasets. w/o Lie: without mapping input and noise into the Lie algebra space, w/o Rot.: without mapping rotation noise into the Lie algebra space, w/o Trans.: without mapping translation noise into the Lie algebra space, w/o CD: without the Chamfer Distance as an additional test metric.

The results (Table 4 and Figure 5) show a significant performance decline when the Lie algebra mapping is removed. This decline underscores the critical role of Lie algebra mapping in linearizing the noise addition and denoising processes. Furthermore, Figure 5 (second column) reveals that the absence of Lie mapping leads to disorganized pose predictions, as the model fails to effectively disentangle pose noise.

The impact of removing rotation noise mapping is particularly pronounced. Without this step, the network struggles to learn pose features due to the inherent nonlinearity and complexity of rotational transformations. The visualization results further confirm this observation, showing erratic and inconsistent pose predictions for parts subjected to rotations.

The removal of translation noise mapping has a relatively smaller impact on the overall performance, though it still introduces noticeable errors in translation estimation. This suggests that while both rotation and translation mappings are important, the rotational component plays a more critical role in ensuring robust pose predictions.

The inclusion of Chamfer Distance loss improves part accuracy by refining the geometric alignment between predicted and ground truth poses. It encourages the model to prioritize local geometric alignment over global pose accuracy.

5 Conclusion and Limitation

This paper presents an SE(3)-equivariant diffusion model for predicting transformed part poses and reassembling fragmented parts into a complete object. The framework consists of three main components: an equivariant feature encoder, a Lie algebra mapping module and an elastic diffusion model. The encoder extracts rotational and translational equivariant features, the Lie algebra mapping enables the model to handle transformations in a linearized manner, and the elastic diffusion model leverages NCSNs with an elastic property, allowing the framework to adapt to varying part quantities. These design elements collectively enhance the model's robustness and adaptability for complex part-based reassembly tasks.

Limitations and future work. The current method is limited to rigid transformations, such as rotation and translation, which restrict its application to scenarios requiring flexible transformations, such as sliding, bending, or stretching. Incorporating flexible transformations without disassembling parts for forward prediction remains a challenging yet essential direction for improvement. Future work will focus on extending the model to handle non-rigid deformations by integrating techniques that can capture local pose changes driven by global structural deformations. This extension would enable the model to predict transformations for objects with complex, flexible morphologies, broadening its applicability to domains such as biomechanical modeling, soft robotics, and human body pose estimation.

Acknowledgements

This work is supported by the National Natural Science Foundation of China under Grant 62471266, as well as Ningbo Major Research and Development Plan under Grant 2023Z224.