SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Shentong Mo<sup>1</sup> Enze Xie<sup>2*</sup> Ruihang Chu<sup>3</sup> Lewei Yao<sup>2</sup>  
Lanqing Hong<sup>2</sup> Matthias Nießner<sup>4</sup> Zhenguo Li<sup>2</sup>
$^{1}$ MBZUAI,  $^{2}$ Huawei Noah's Ark Lab,  $^{3}$ CUHK,  $^{4}$ TUM
https://DiT-3D.github.io
Recent Diffusion Transformers (e.g. DiT [1]) have demonstrated their powerful effectiveness in generating high-quality 2D images. However, it is still being determined whether the Transformer architecture performs equally well in 3D shape generation, as previous 3D diffusion methods mostly adopted the U-Net architecture. To bridge this gap, we propose a novel Diffusion Transformer for 3D shape generation, namely DiT-3D, which can directly operate the denoising process on voxelized point clouds using plain Transformers. Compared to existing U-Net approaches, our DiT-3D is more scalable in model size and produces much higher quality generations. Specifically, the DiT-3D adopts the design philosophy of DiT [1] but modifies it by incorporating 3D positional and patch embeddings to adaptively aggregate input from voxelized point clouds. To reduce the computational cost of self-attention in 3D shape generation, we incorporate 3D window attention into Transformer blocks, as the increased 3D token length resulting from the additional dimension of voxels can lead to high computation. Finally, linear and devoxelization layers are used to predict the denoised point clouds. In addition, our transformer architecture supports efficient fine-tuning from 2D to 3D, where the pre-trained DiT-2D checkpoint on ImageNet can significantly improve DiT-3D on ShapeNet. Experimental results on the ShapeNet dataset demonstrate that the proposed DiT-3D achieves state-of-the-art performance in high-fidelity and diverse 3D point cloud generation. In particular, our DiT-3D decreases the 1-Nearest Neighbor Accuracy of the state-of-the-art method by 4.59 and increases the Coverage metric by 3.51 when evaluated on Chamfer Distance.
In recent times, there has been a growing interest in exploring the potential of diffusion transformers for high-fidelity image generation, as evinced by a series of scholarly works [1, 2, 3, 4]. Notably, a seminal work by Peebles et al. [1] proposed the replacement of the widely-used U-Net backbone with a scalable transformer. Specifically, the proposed method operates on latent patches by training latent 2D diffusion models. However, the efficacy of plain diffusion transformers for 3D shape generation has yet to be explored, as most existing 3D diffusion approaches continue to adopt the U-Net backbone.
Generating high-fidelity point clouds for 3D shape generation is a challenging and significant problem. Early generative methods [5, 6, 7] addressed this problem by directly optimizing heuristic loss objectives, such as Chamfer Distance (CD) and Earth Mover's Distance (EMD). More recent works [8, 9, 10, 11] have explored the usage of the generative adversarial network (GAN)-based
Figure 1: Examples of high-fidelity and diverse 3D point clouds produced from DiT-3D.
and flow-based models to generate 3D point clouds from a probabilistic perspective. Recently, researchers [12, 13, 14, 15] have turned to various denoising diffusion probabilistic models (DDPMs) to generate entire shapes from random noise. For instance, PVD [12] employed the point-voxel representation of 3D shapes as input to DDPMs. They reversed the diffusion process from observed point clouds to Gaussian noise by optimizing a variational lower bound to the likelihood function. Recently, the Diffusion Transformer (DiT) [1, 2] has been shown to surpass the U-Net architecture in 2D image generation, owing to its simple design and superior generative performance. Consequently, we investigate the potential of the Diffusion Transformer for 3D generation. However, extending the 2D DiT to 3D poses two significant challenges: (1) Point clouds are intrinsically unordered, unlike images where pixels are ordered; and (2) The tokens in 3D space have an additional dimension compared to 2D images, resulting in a substantial increase in computational cost.
This work introduces DiT-3D, a novel diffusion transformer architecture designed for 3D shape generation that leverages the denoising process of DDPM on 3D point clouds. The proposed model inherits the simple design of the modules in DiT-2D, with only minor adaptations to enable it to generalize to 3D generation tasks. To tackle the challenge posed by the unordered data structure of point clouds, we convert the point cloud into a voxel representation. DiT-3D employs 3D positional embedding and 3D patch embedding on the voxelized point clouds to extract point-voxel features and effectively process the unordered data. Furthermore, to address the computational cost associated with a large number of tokens in 3D space, we introduce a 3D window attention operator instead of the vanilla global attention in DiT-2D. This operator significantly reduces training time and memory usage, making DiT-3D feasible for large-scale 3D generation tasks. Finally, we utilize linear and devoxelization layers to predict the noised point clouds in the reversed process to generate final 3D shapes.
In order to address the computational cost associated with a large number of tokens in 3D space, we also introduce a parameter-efficient tuning method to utilize the pre-trained DiT-2D model on ImageNet as initialization for DiT-3D (window attention shares the same parameters with vanilla attention). Benefiting from the substantial similarity between the network structure and parameters of DiT-3D and DiT-2D, the representations learned on ImageNet significantly improve 3D generation, despite the significant domain disparity between 2D images and 3D point clouds. To our knowledge, we are the first to achieve parameter-efficient fine-tuning from 2D ImageNet pre-trained weights for high-fidelity and diverse 3D shape generation. In particular, we highly decrease the training parameters from 32.8MB to only 0.09MB.
We present a comprehensive evaluation of DiT-3D on a diverse set of object classes in the ShapeNet benchmark, where it achieves state-of-the-art performance compared to previous non-DDPM and DDPM-based 3D shape generation methods. Qualitative visualizations further emphasize the efficacy of DiT-3D in generating high-fidelity 3D shapes. Extensive ablation studies confirm the significance of 3D positional embeddings, window attention, and 2D pre-training in 3D shape generation. Moreover, we demonstrate that DiT-3D is easily scalable regarding patch sizes, voxel sizes, and model sizes. Our
findings align with those of DiT-2D, where increasing the model size leads to continuous performance improvements. In addition, our parameter-efficient fine-tuning from DiT-2D ImageNet pre-trained weights highly decreases the training parameters while achieving competitive generation performance. By only training 0.09MB parameters of models from the source class to the target class, we also achieve comparable results of quality and diversity in terms of all metrics.
- We present DiT-3D, the first plain diffusion transformer architecture for point cloud shape generation that can effectively perform denoising operations on voxelized point clouds.
- We make several simple yet effective modifications on DiT-3D, including 3D positional and patch embeddings, 3D window attention, and 2D pre-training on ImageNet. These modifications significantly improve the performance of DiT-3D while maintaining efficiency.
- Extensive experiments on the ShapeNet dataset demonstrate the state-of-the-art superiority of DiT-3D over previous non-DDPM and DDPM baselines in generating high-fidelity shapes.
3D Shape Generation. 3D shape generation aims to synthesize high-fidelity point clouds or meshes using generative models, such as variational autoencoders [16, 17, 18], generative adversarial networks [19, 8, 20], and normalized flows [9, 10, 11]. Typically, PointFlow [9] utilized a probabilistic framework based on the continuous normalizing flow to generate 3D point clouds from two-level hierarchical distributions. ShapeGF [21] trained a score-matching energy-based network to learn the distribution of points across gradient fields using Langevin dynamics. More recently, GET3D [14] leveraged a signed distance field (SDF) and a texture field as two latent codes to learn a generative model that directly generates 3D meshes. In this work, we mainly focus on denoising diffusion probabilistic models for generating high-fidelity 3D point clouds starting from random noise, where point and shape distributions are not separated.
Diffusion Models. Diffusion models [22, 23, 24] have been demonstrated to be effective in many generative tasks, such as image generation [25], image restoration [26], speech generation [27], and video generation [28]. Denoising diffusion probabilistic models (DDPMs) [22, 23] utilized a forward noisin process that gradually adds Gaussian noise to images and trained a reverse process that inverts the forward process. In recent years, researchers [29, 12, 13, 30, 15, 31, 32] have tried to explore diverse pipelines based on diffusion probabilistic models to achieve 3D shape generation. For example, PVD [12] applied DDPM based on PVCNNs [33] on the point-voxel representation of 3D shapes with structured locality into point clouds. To improve the generation quality, LION [13] used two DDPMs to learn a hierarchical latent space based on a global shape latent representation and a point-structured latent space separately. Different from them, we will solve the 3D shape generation problem in our approach by designing a plain transformer-based architecture backbone to replace the U-Net backbone for reversing the diffusion process from observed point clouds to Gaussian noise. Meanwhile, our 3D plain diffusion transformer supports multi-class training with learnable class embeddings as the condition and parameter-efficient fine-tuning with modality and domain transferability differ from DDPM-based 3D generation approaches discussed above.
Transformers in Diffusion Generation. Diffusion Transformers [1, 2, 3, 4] have recently shown their impressive capacity to generate high-fidelity images. For instance, Diffusion Transformer (DiT) [1] proposed a plain diffusion Transformer architecture to learn the denoising diffusion process on latent patches from a pre-trained pre-trained variational autoencoder model in Stable Diffusion [34]. U-ViT [2] incorporated all the time, condition, and noisy image patches as tokens and utilized a Vision transformer(ViT) [35]-based architecture with long skip connections between shallow and deep layers. More recently, UniDiffuser [3] designed a unified transformer for diffusion models to handle input types of different modalities by learning all distributions simultaneously. While those diffusion transformer approaches achieve promising performance in 2D image generation, how a plain diffusion transformer performs on 3D shape generation is still being determined. In contrast, we develop a novel plain diffusion transformer for 3D shape generation that can effectively perform denoising operations on voxelized point clouds. Furthermore, the proposed DiT-3D can support parameter-efficient fine-tuning with transferability across modality and domain.
Figure 2: Illustration of the proposed Diffusion Transformers (DiT-3D) for 3D shape generation. The plain diffusion transformer takes voxelized point clouds as input, and a patchification operator is used to generate token-level patch embeddings, where 3D positional embeddings are added together. Then, multiple transformer blocks based on 3D window attention extract point-voxel representations from all input tokens. Finally, the unpatchified voxel tensor output from a linear layer is devoxelized to predict the noise in the point cloud space.
Given a set of 3D point clouds, we aim to learn a plain diffusion transformer for synthesizing new high-fidelity point clouds. We propose a novel diffusion transformer that operates the denoising process of DDPM on voxelized point clouds, namely DiT-3D, which consists of two main modules: Design DiT for 3D Point Cloud Generation in Section 3.2 and Efficient Modality/Domain Transfer with Parameter-efficient Fine-tuning in Section 3.3.
In this section, we first describe the problem setup and notations and then revisit denoising diffusion probabilistic models (DDPMs) for 3D shape generation and diffusion transformers on 2D images.
Problem Setup and Notations. Given a set  $S = \{\mathbf{p}_i\}_{i=1}^S$  of 3D shapes with  $M$  classes, our goal is to train a plain diffusion transformer from these point clouds for generating high-fidelity point clouds. For each point cloud  $\mathbf{p}_i$ , we have  $N$  points for  $x, y, z$  coordinates, that is  $\mathbf{p}_i \in \mathbb{R}^{N \times 3}$ . Note that we have a class label for the 3D shape  $\mathbf{p}_i$ , which is denoted as  $\{y_i\}_{i=1}^M$  with  $y_i$  for the ground-truth category entry  $i$  as 1. During the training, we take the class label as input to achieve classifier-free guidance in conditional diffusion models, following the prior diffusion transformer (i.e., DiT [1]) on images.
Revisit DDPMs on 3D Shape Generation. To solve the 3D shape generation problem, previous work [12] based on denoising diffusion probabilistic models (DDPMs) define a forward noisng process that gradually applies noise to real data  $\mathbf{x}_0$  as  $q(\mathbf{x}_t|\mathbf{x}_{t - 1}) = \mathcal{N}(\mathbf{x}_t;\sqrt{1 - \beta_t}\mathbf{x}_{t - 1},\beta_t\mathbf{I})$  where  $\beta_{t}$  is a Gaussian noise value between 0 and 1. In particular, the denoising process produces a series of shape variables with decreasing levels of noise, denoted as  $\mathbf{x}_T,\mathbf{x}_{T - 1},\dots ,\mathbf{x}_0$  , where  $\mathbf{x}_T$  is sampled from a Gaussian prior and  $\mathbf{x}_0$  is the final output. With the reparameterization trick, we can have  $\mathbf{x}_t = \sqrt{\bar{\alpha}_t}\mathbf{x}_0 + \sqrt{1 - \bar{\alpha}_t}\pmb {\epsilon}$  , where  $\pmb {\epsilon}\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ $\alpha_{t} = 1 - \beta_{t}$  , and  $\bar{\alpha}_{t} = \prod_{i = 1}^{t}\alpha_{i}$
For the reverse process, diffusion models are trained to learn a denoising network  $\theta$  for inverting forward process corruption as  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t) = \mathcal{N}(\mathbf{x}_{t - 1};\boldsymbol {\mu}_\theta (\mathbf{x}_t,t),\sigma_t^2\mathbf{I})$ . The training objective is to maximize a variational lower bound of the negative log data likelihood that involves all of  $\mathbf{x}_0,\dots ,\mathbf{x}_T$
$$
\mathcal {L} = \sum_ {t} - p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {0} \mid \mathbf {x} _ {1}\right) + \mathcal {D} _ {\mathrm {K L}} \left(q \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}, \mathbf {x} _ {0}\right) \mid \mid p _ {\boldsymbol {\theta}} \left(\mathbf {x} _ {t - 1} \mid \mathbf {x} _ {t}\right)\right)) \tag {1}
$$
where  $\mathcal{D}_{\mathrm{KL}}(\cdot ||\cdot)$  denotes the KL divergence measuring the distance between two distributions. Since both  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t))$  and  $q(\mathbf{x}_{t - 1}|\mathbf{x}_t,\mathbf{x}_0)$  are Gaussians, we can reparameterize  $\mu_{\theta}(\mathbf{x}_t,t)$  to predict the noise  $\epsilon_{\theta}(\mathbf{x}_t,t)$ . In the end, the training objective can be reduced to a simple mean-squared loss between the model output  $\epsilon_{\theta}(\mathbf{x}_t,t)$  and the ground truth Gaussian noise  $\epsilon$  as:  $\mathcal{L}_{\mathrm{simple}} = \| \epsilon - \epsilon_{\theta}(\mathbf{x}_t,t)\|^2$ . After  $p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t))$  is trained, new point clouds can be generated by progressively sampling  $\mathbf{x}_{t - 1}\sim p_{\theta}(\mathbf{x}_{t - 1}|\mathbf{x}_t))$  by using the reparameterization trick with initialization of  $\mathbf{x}_T\sim \mathcal{N}(\mathbf{0},\mathbf{I})$ .
Revisit Diffusion Transformer (DiT) on 2D Image Generation. To generate high-fidelity 2D images, DiT proposed to train latent diffusion models (LDMs) with Transformers as the backbone, consisting of two training models. They first extract the latent code  $\mathbf{z}$  from an image sample  $\mathbf{x}$  using an autoencoder with an encoder  $f_{\mathrm{enc}}(\cdot)$  and a decoder  $f_{\mathrm{dec}}(\cdot)$ , that is,  $\mathbf{z} = f_{\mathrm{enc}}(\mathbf{x})$ . The decoder is used to reconstruct the image sample  $\hat{\mathbf{x}}$  from the latent code  $\mathbf{z}$ , i.e.,  $\hat{\mathbf{x}} = f_{\mathrm{dec}}(\mathbf{z})$ . Based on latent codes  $\mathbf{z}$ , a latent diffusion transformer with multiple designed blocks is trained with time embedding  $\mathbf{t}$  and class embedding  $\mathbf{c}$ , where a self-attention and a feed-forward module are involved in each block. Note that they apply patchification on latent code  $\mathbf{z}$  to extract a sequence of patch embeddings and depatchification operators are used to predict the denoised latent code  $\mathbf{z}$ .
Although DDPMs achieved promising performance on 3D shape generation, they can only handle single-class training based on PVCNNs [33] as the encoder to extract 3D representations, and they cannot learn explicit class-conditional embeddings. Furthermore, we are not able to directly transfer their single-class pre-trained model to new classes with parameter-efficient fine-tuning. Meanwhile, we empirically observe that the direct extension of DiT [1] on point clouds does not work. To address this problem, we propose a novel plain diffusion transformer for 3D shape generation that can effectively achieve the denoising processes on voxelized point clouds, as illustrated in Figure 2.
3.2 Diffusion Transformer for 3D Point Cloud Generation
To enable denoising operations using a plain diffusion transformer, we propose several adaptations to 3D point cloud generation in Figure 2 within the framework of DiT [1]. Specifically, our DiT-3D model accepts voxelized point clouds as input and employs a patchification operator to generate token-level patch embeddings. We add 3D positional embeddings to these embeddings and extract point-voxel representations from all input tokens using multiple transformer blocks based on 3D window attention. Finally, we apply a devoxelized linear layer to the unpatchified voxel output, allowing us to predict the noise in the point cloud space.
Denoising on Voxelized Point Clouds. Point clouds are inherently unordered, unlike images where pixels follow a specific order. We encountered difficulty in our attempt to train a diffusion transformer on point coordinates due to the sparse distribution of points in the 3D embedding space. To address this issue, we decided to voxelize the point clouds into dense representations, allowing the diffusion transformers to extract point-voxel features. Our approach differs from DiT [1], which utilizes latent codes  $\mathbf{z}$  to train the latent diffusion transformer. Instead, we directly train the denoising process on voxelized point clouds using the diffusion transformer. For each point cloud  $\mathbf{p}_i\in \mathbb{R}^{N\times 3}$  with  $N$  points for  $x,y,z$  coordinates, we first voxelize it as input  $\mathbf{v}_i\in \mathbb{R}^{V\times V\times V\times 3}$ .
3D Positional and Patch Embeddings. With the voxel input  $\mathbf{v}_i\in \mathbb{R}^{V\times V\times V\times 3}$ , we introduce patchification operator with a patch size  $p\times p\times p$  to generate a sequence of patch tokens  $\mathbf{t}\in \mathbb{R}^{L\times 3}$ .  $L = (V / p)^{3}$  denotes the total number of patched tokens. A 3D convolution layer is applied on patch tokens to extract patch embeddings  $\mathbf{e}\in \mathbb{R}^{L\times D}$ , where  $D$  is the dimension of embeddings. To adapt to our voxelized point clouds, we add frequency-based sine-cosine 3D positional embeddings instead of the 2D version in DiT [1] to all input tokens. Based on these patch-level tokens, we introduce time embeddings  $\mathbf{t}$  and class embeddings  $\mathbf{c}$  as input to achieve multi-class training with learnable class embeddings as the condition, which differs from existing 3D generation approaches with U-Net as the backbone.
3D Window Attention. Due to the increased token length resulting from the additional dimension in 3D space, the computational cost of 3D Transformers can be significantly high. To address this issue, we introduce efficient 3D window attention into Transformer blocks blocks to propagate point-voxel
features in efficient memory usage. For the original multi-head self-attention process with each of the heads  $Q, K, V$  have the same dimensions  $L \times D$ , where  $L = (V / p)^3$  is the length of input tokens, we can have the attention operator as:
$$
\operatorname {A t t e n t i o n} (Q, K, V) = \operatorname {S o f t m a x} \left(\frac {Q K ^ {\top}}{\sqrt {D _ {h}}} V\right) \tag {2}
$$
where  $D_h$  is the dimension size of each head. The computational complexity of this process is  $\mathcal{O}(L^2)$ , which will be largely expensive for high voxel resolutions. Inspired by [36], we extend the 2D window attention operator to a 3D one for 3D input tokens instead of vanilla global attention. This process uses a window size of  $R$  to reduce the length of total input tokens as
$$
\hat {K} = \operatorname {R e s h a p e} \left(\frac {L}{R ^ {3}}, D \cdot R ^ {3}\right) \tag {3}
$$
$$
K = \operatorname {L i n e a r} (D \cdot R ^ {3}, D) (\hat {K})
$$
where  $K$  is the input tokens to be reduced. Reshape  $\left(\frac{L}{R^3}, D \cdot R^3\right)$  denotes to reshape  $K$  to the one with shape of  $\frac{L}{R^3} \times (D \cdot R^3)$ , and Linear $(C_{in}, C_{out})(\cdot)$  denotes to a linear layer with a  $C_{in}$ -dimensional tensor as input and a  $C_{out}$ -dimensional tensor as output. Therefore, the new  $K$  has the shape of  $\frac{L}{R^3} \times D$ . As a result, the complexity of the self-attention operator in Equation (2) is reduced from  $\mathcal{O}(L^2)$  to  $\mathcal{O}\left(\frac{L^2}{R^3}\right)$ . In our experiments, we set  $R$  to 4 in the default setting.
Devoxelized Prediction. Since the transformers blocks are implemented on voxelized point clouds, we can not directly use a standard linear decoder to predict the output noise  $\epsilon_{\theta}(\mathbf{x}_t,t)$  from point clouds. In order to generate the output noise, we devoxelize output tokens from the linear decoder. We first apply the final layer norm and linearly decode each token into a  $p\times p\times p\times L\times 3$  tensor, where  $L$  is the total number of input tokens. Then we unpatchify the decoded token into a voxel tensor with the shape of  $V\times V\times V\times 3$ . Finally, the unpatchified voxel tensor is devoxelized into a  $N\times 3$  tensor as the output noise  $\epsilon_{\theta}(\mathbf{x}_t,t)$ , matching with the ground truth Gaussian noise  $\epsilon$  in the point cloud space.
Model Scaling. Our DiT-3D is designed to be scalable, adapting to varying voxel sizes, patch sizes, and model sizes. Specifically, it can flexibly accommodate voxel dimensions of 16, 32, 64, patch dimensions of 2, 4, 8, and model complexity ranging from Small, Base, Large and Extra Large, as demonstrated in DiT [1]. For instance, a model designated as DiT-3D-S/4 refers that it utilizes the Small configuration of the DiT model [1], with a patch size  $p$  of 4.
3.3 Efficient Modality/Domain Transfer with Parameter-efficient Fine-tuning
Leveraging the scalability of the plain diffusion transformer, we investigate parameter-efficient fine-tuning for achieving modality and domain transferability. To facilitate modality transfer from 2D to 3D, we can leverage the knowledge pre-trained on large-scale 2D images using DiT [1]. For domain transfer from a source class to target classes, we train DiT-3D on a single class (e.g. chair) and transfer the model's parameters to other classes (e.g. airplane, car).