DiT-3D: Исследование простых диффузионных трансформеров для генерации 3D-форм

DiT-3D: Exploring Plain Diffusion Transformers for 3D Shape Generation

July 4, 2023

Авторы: Shentong Mo, Enze Xie, Ruihang Chu, Lewei Yao, Lanqing Hong, Matthias Nießner, Zhenguo Li

cs.AI

Аннотация

Недавние диффузионные трансформеры (например, DiT) продемонстрировали свою мощную эффективность в генерации высококачественных 2D-изображений. Однако до сих пор неясно, насколько хорошо архитектура трансформеров справляется с генерацией 3D-форм, поскольку предыдущие методы 3D-диффузии в основном использовали архитектуру U-Net. Чтобы устранить этот пробел, мы предлагаем новый диффузионный трансформер для генерации 3D-форм, а именно DiT-3D, который может напрямую выполнять процесс удаления шума на вокселизированных облаках точек с использованием простых трансформеров. По сравнению с существующими подходами на основе U-Net, наш DiT-3D более масштабируем по размеру модели и генерирует результаты значительно более высокого качества. В частности, DiT-3D использует философию проектирования DiT, но модифицирует её, добавляя 3D-позиционные и патч-эмбеддинги для адаптивной агрегации входных данных из вокселизированных облаков точек. Чтобы снизить вычислительные затраты на само-внимание в генерации 3D-форм, мы внедряем 3D-оконное внимание в блоки трансформеров, так как увеличение длины 3D-токенов из-за дополнительного измерения вокселей может привести к высоким вычислительным затратам. Наконец, линейные и девокселизационные слои используются для предсказания очищенных от шума облаков точек. Кроме того, наша архитектура трансформеров поддерживает эффективную тонкую настройку с 2D на 3D, где предварительно обученная модель DiT-2D на ImageNet может значительно улучшить DiT-3D на ShapeNet. Экспериментальные результаты на наборе данных ShapeNet демонстрируют, что предложенный DiT-3D достигает наилучших показателей в генерации высококачественных и разнообразных 3D-облаков точек. В частности, наш DiT-3D снижает метрику 1-Nearest Neighbor Accuracy у лучшего существующего метода на 4,59 и увеличивает метрику Coverage на 3,51 при оценке по расстоянию Чамфера.

English

Recent Diffusion Transformers (e.g., DiT) 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 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.