Part-aware Shape Generation with Latent 3D Diffusion of Neural Voxel Fields

Inspired by the preliminary work, we introduce a latent 3D diffusion model for the generation of high-resolution neural voxel fields. Our model offers two key advantages in comparison to previous diffusion-based methods. Unlike conventional 2D diffusion models [21, 50, 34, 51], we utilize 3D convolution to directly denoise the voxel field in 3D space, which enhances our modeling capabilities in terms of 3D geometry and consistency. On the other hand, due to the substantial computational demands of 3D convolution, training a 3D diffusion model with high-resolution voxel fields proves challenging, thereby constraining the generative quality. To address this limitation, we shift the denoising process to a latent space with reduced resolution, effectively eliminating constraints on voxel field resolution.

As shown in Fig. 2, the proposed latent 3D diffusion model comprises a 3D autoencoder and a 3D denoising UNet. The 3D autoencoder has an encoder for compressing the voxel field $\mathbf{V}^{\boldsymbol{f}}_{0}$ to the latent code $\mathbf{L}_{0}$ and a decoder for recovering the original voxel field from $\mathbf{L}_{0}$, i.e. $\mathbf{L}_{0}=\mathcal{E}(\mathbf{V}^{\boldsymbol{f}}_{0}),\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}=\mathcal{D}(\mathbf{L}_{0})$. Here, $\mathcal{E}$ and $\mathcal{D}$ represent the encoder and decoder of the 3D autoencoder; $\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}$ is the reconstructed voxel field. In practice, the resolution of $\mathbf{L}_{0}$ is 64$\times$ lower than $\mathbf{V}^{\boldsymbol{f}}_{0}$. After obtaining the latent code, we adopt a 3D UNet architecture to conduct the diffusion process in the latent space. Similar to Eq. 3, we sample the noisy code $\mathbf{L}_{t}$ by adding the normal noise $\boldsymbol{\epsilon}$ to $\mathbf{L}_{0}$. Then, the noisy code is fed into the 3D UNet with the time $t$ together, outputting the estimations of $\mathbf{L}_{0}$ and $\boldsymbol{\epsilon}$. This process is formulated by:

where $\boldsymbol{\phi}$ is the parameter of 3D UNet, $\mathbf{L}_{\boldsymbol{\phi}}$ and $\boldsymbol{\epsilon}_{\boldsymbol{\phi}}$ are the predictions of $\mathbf{L}_{0}$ and $\boldsymbol{\epsilon}$. Note that we use the special form of the attenuation function: $\boldsymbol{a}_{t}=-\mathbf{L}_{0}$.

For conditioned generation, we utilize the classifier-free [52] manner to synthesize shapes corresponding to the conditioned input. Specifically, the conditioned input is fed into a feature extractor, resulting in a feature embedding. We adopt the feature embedding as an additional input of the 3D UNet, thus we obtain the $\mathbf{L}_{\boldsymbol{\phi}}$ and $\boldsymbol{\epsilon}_{\boldsymbol{\phi}}$ by: $\mathbf{L}_{\boldsymbol{\phi}},\boldsymbol{\epsilon}_{\boldsymbol{\phi}}=\text{UNet}_{\boldsymbol{\phi}}(\mathbf{L}_{t},t,\mathbf{e})$, where $\mathbf{e}$ is the feature embedding of the conditioned input.

3D UNet vs. 2D UNet. Actually, though the latent code $\mathbf{L}_{t}$ is a 4D tensor, it is possible to apply the 2D convolutional operator to this type of data. In practice, we can reshape $\mathbf{L}_{t}$ into a 3D tensor, i.e. $\text{reshape}(\mathbb{R}^{C\times X\times Y\times Z})\rightarrow\mathbb{R}^{CX\times Y\times Z}$, thus adopting 2D UNet to model the denoising process. However, we found that the loss curve did not converge and it failed to obtain satisfactory results (details can be found in Sec. IV-C). We assume that the 2D UNet architecture hardly learns the 3D consistency of the voxel field, resulting in a perturbed result. In contrast, the 3D UNet architecture provides reasonable performance, demonstrating the necessity of adopting 3D UNet for denoising the neural voxel field.

After the denoising process of the latent 3D diffusion model, we can have a neural voxel field containing the density and color features. However, it can not provide the structural information and rendering results. Therefore, we propose a part-aware shape decoder to generate part-aware details and rendering results.

As shown in Fig. 2, to model the part-aware information, the shape decoder integrates the part code into the voxel field and utilizes the part attention to build the relations between different parts of the voxel field. Specifically, the part code is a learnable embedding and can be represented by $\boldsymbol{z}\in\mathbb{R}^{K\times D}$, where $K$ is the number of parts, $D$ is the dimension of the embedding. To incorporate the part code into the neural field, we first project the reconstructed voxel field $\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}$ to have the same dimension as $\boldsymbol{z}$, and then flatten $\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}$ to be a 2D tensor like $\mathbb{R}^{XYZ\times D}$. Thus, we can utilize the multi-head cross-attention [53] to build a strong dependency between the part code and the voxel field. Moreover, to divide the voxel field accurately, we employ a self-attention to model the long-range relations between the different parts of the voxel field. After the cross-attention and self-attention, we reshape the tensor back to be a 4D voxel field. We formulate this process by the following equation:

where MHSA and MHCA represent the multi-head self-attention and cross-attention.

Incorporated with the part code, the refined voxel field $\overline{\mathbf{V}}^{\boldsymbol{f}}_{0}$ contains both the color and structural information. In this way, we can utilize volume rendering to generate images and part maps of arbitrary views. Given a 3D position $\boldsymbol{x}$, we can query the point feature of the voxel field, and acquire the color emission and part class probability with MLP layers.

where $\boldsymbol{p}\in\mathbb{R}^{K}$ is the part class probability. In contrast to Eq. 3, we incorporate the part code into the MLP layers collectively to encode the structural information into individual points, thereby enhancing the generative intricacies. Subsequently, we derive the RGB image and semantic part map via ray casting following Eq. 2.

Several methods [54, 55] in the literature have combined the diffusion model with the neural field decoder, typically training these two components independently. However, independent training neglects the consideration of noise in the prior term $\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}$ during the training of the neural field decoder, leading to biased learning of the decoder. In contrast, we propose an end-to-end training approach for the latent 3D diffusion model and the part-aware shape decoder, thereby avoiding noisy generation.

Gradient Skip. The diffusion model is trained on the latent space while the part-aware shape decoder is applied to the voxel field directly. Besides, there is the decoder $\mathcal{D}$ of the 3D autoencoder between the two modules connecting the latent space and the voxel field space. Obviously, training the two modules jointly needs the backward propagation of the gradient of $\mathcal{D}$. However, this implementation incurs significant computational costs and hinders the effective feedback of gradients. To mitigate this problem, we propose to skip the gradient of $\mathcal{D}$, allowing the gradient of the part-aware shape decoder to directly propagate back to the latent space. Concretely, according to the Chain Rule, we have:

where $\mathcal{L}_{rend}$ is the rendering loss. We ignore the term $\partial\hat{\mathbf{V}}^{\boldsymbol{f}}_{0}/\partial\mathbf{L}_{\boldsymbol{\phi}}$ to skip the gradient of the decoder in the 3D autoencoder, resulting in an efficient formulation of gradient calculation:

The main training objective is the summation of the diffusion loss and the rendering loss.

where $\hat{\boldsymbol{c}}(\boldsymbol{r})$ and $\hat{\boldsymbol{p}}(\boldsymbol{r})$ are the rendered RGB image and part map from ray $\boldsymbol{r}$, and $\boldsymbol{g}_{\boldsymbol{c}}(r)$ and $\boldsymbol{g}_{\boldsymbol{p}}(r)$ are the corresponding ground truths; CE means the cross-entropy loss function. Since the predicted voxel field of the diffusion model is reasonable only with the time $t$ close to zero, we attach a time-dependent weight $\lambda_{t}$ to the rendering loss: $\lambda_{t}=-\log(t)$.

Except for the main training objective, we adopt two additional regularization loss functions. First, we conduct the total variation on the neural voxel fields to restrain the noise generating:

where $\Delta_{x}^{2}(\hat{\mathbf{V}}^{\boldsymbol{f}}_{0},dim)$ denotes the squared difference between the value of $dim$-th channel in voxel: $v:=(i;j;k)$ and the $dim$-th value in voxel $(i+1;j;k)$, which can be extended to $\Delta_{y}^{2}(\hat{\mathbf{V}}^{\boldsymbol{f}}_{0},dim)$ and $\Delta_{z}^{2}(\hat{\mathbf{V}}^{\boldsymbol{f}}_{0},dim)$. Moreover, we constrain the value of part code $\boldsymbol{z}$ via its L2 norm $\|\boldsymbol{z}\|_{2}$ attached with the weight 0.0001.

Quantitative comparison against state of the art. In Tab. I, Tab. II and Tab. III, We have conducted a quantitative comparison between our proposed method and state-of-the-art approaches. Tab. I reports the metrics of generated image quality, while Tab. II and Tab. III report the geometry metrics calculated by Chamfer distance and Earth Mover’s distance. EG3D [13] presents a GAN-based framework that introduces a sophisticated hybrid explicit-implicit network architecture for synthesizing the triplane representation. Despite its lightweight nature, the triplane representation struggles to encompass intricate details, resulting in unsatisfied rendering quality and geometry. Part-NeRF [17] employs multiple neural radiance fields to represent a single object, but the complex and redundant architecture increases the learning difficulty, ultimately yielding lower-quality results. DiffRF [18] also utilizes the diffusion model in 3D space; however, it directly applies the 3D UNet to the voxel field. This implementation imposes significant computational demands, constraining the resolution of the voxel fields and leading to sub-optimal performances. Leveraging the 3D latent diffusion design enables the generation of 3D shapes with enhanced geometric details, leading to the attainment of the highest MMD and COV metrics across all classes. Additionally, the part-aware shape decoder facilitates the establishment of relationships between distinct parts, thereby improving the rendering of different components and resulting in the best FID metric across all classes.

Qualitative comparison against state of the art. As depicted in Fig. 3, we plot the visual generations of various state-of-the-art methods. Leveraging the latent 3D diffusion model, we can generate high-resolution voxel fields, resulting in superior rendering quality. Similar to ours, Part-NeRF also produces part-aware rendering results. However, it employs an unsupervised approach that leads to the synthesis of meaningless structures. To ensure a fair comparison, we employ semantic part maps to supervise Part-NeRF, denoting this as Part-NeRF-S. As shown in Fig. 3, benefiting from the design of latent diffusion and part-aware shape decoder, our approach still outperforms Part-NeRF-S in both texture and part quality. Both quantitative and qualitative results indicate that the proposed method can generate accurate part-aware shapes and render high-quality images.

We ablate different components of the proposed method to show their importance.

Importance of the latent 3D diffusion. Different from previous methods, this paper proposes the latent 3D diffusion that transfers the diffusion process into latent space, increasing the generative resolution significantly and preserving the high-quality generation. We further apply the diffusion process to the voxel field directly to observe the performance change. Due to the limitation of computational cost, we set the resolution of voxel fields to 32$\times$32$\times$32. Tab. IV reports the performance comparison. Without the latent diffusion, the FID drops a lot, and the geometry metric of MMD and COV drop 0.36 and 4.9, which shows the importance of latent diffusion. Moreover, we visualize the qualitative results of the two different settings in Fig. 5, and the results reveal that the latent 3D diffusion brings high-quality geometry and rendering details.

Importance of the part-aware shape decoder. The part-aware shape decoder incorporates the part code into the neural voxel field, generating the structural shape with accurate rendering results. To demonstrate its importance, we replace it with two independent MLPs, rendering the color and part maps from the voxel field directly. Tab. V shows that the part-aware shape decoder can bring more improvements in rendering quality. As depicted in Fig. 4, the part-aware shape decoder helps the model learn the accurate structural information, which allows for consistent rendering results of the local region.

Importance of the gradient skip. The gradient skip technique aims to reduce redundant gradient calculations and computational costs by jointly learning the latent 3D diffusion model and the part-aware shape decoder. In Fig. 11, we compare our method with the split training strategy, which does not utilize the gradient skip technique. It is evident from the comparison that the split training strategy lacking gradient skipping leads to the learning of a biased neural voxel field, consequently yielding inferior rendered images.

Importance of 3D UNet. We employ a 3D UNet as the network architecture to conduct the denoising process, and we have analyzed that the denoising network can be replaced by a 2D UNet possibly. However, we found that the 2D UNet failed to generate meaningful results as shown in Fig. 9. Compared to 3D UNet, the training loss curve of 2D UNet is not converged, culminating in a noisy generation. This observation indicates that it is necessary to utilize the 3D convolutional operator to extract features for the 3D voxel fields.

Effect of sampling steps. By leveraging rapid sampling from the decoupled diffusion model, we achieve satisfactory results in just 30 steps. Tab. VI presents a comparison of performance based on various sampling steps for chair generation. With an increase in the sampling step, performance shows a gradual improvement. Further increasing the sampling steps results in only a marginal improvement, suggesting that the generation quality is nearing saturation. It is observed that the sampling step has a greater impact on FID and KID metrics, highlighting its influence on image quality rather than geometry.

Following the 2D diffusion model [50], we can utilize two different voxel fields to generate intermediate shapes with spherical linear interpolation [61]. As shown in Fig. 7, our model can achieve relatively smooth interpolation between the two shapes. In contrast, the model without the part-aware shape decoder generates unreasonable contents, which further indicates the effectiveness of the part-aware shape decoder. With the part-aware shape decoder, the generated intermediate shapes consistently maintain a precise structure, suggesting that the decoder effectively learns to decompose the neural voxel fields with the guidance of the part code.

Single image to shape can be seen as a conditional generation task in which we use an image as conditional input and generate a shape related to the image. In practice, we extract the multi-scale image features via the Swin backbone network, and then concatenate the obtained features with the multi-scale features of the denoising UNet. As shown in Fig. 6, given a reference image as a condition, our model can generate reasonable shapes with consistent geometry and accurate structure.

Text-guided shape generation is a form of conditioned generative method, where the text is embedded into a feature to control the generation process. We perform text-guided shape generation on Chairs and Tables. As shown in Fig. 8, our model is capable of generating reasonable shapes based on the provided texts. In fact, the textual condition provides a relatively coarse level of control for generation, resulting in better geometric accuracy compared to texture detail.

Due to the part-aware nature of the generated shapes, it becomes feasible to combine parts from different shapes to create new designs. Notably, these shapes conform to a shared coordinate system, enabling the direct integration of disparate parts to create a cohesive and compact shape. Fig. 10 illustrates the visual results of shape mixing on four different object categories. Our approach excels in producing clear part shapes, facilitating the composition of diverse shapes.