Snelle training van Diffusion Transformer met extreme masking voor generatie van 3D-puntenwolken

Fast Training of Diffusion Transformer with Extreme Masking for 3D Point Clouds Generation

December 12, 2023

Auteurs: Shentong Mo, Enze Xie, Yue Wu, Junsong Chen, Matthias Nießner, Zhenguo Li

cs.AI

Samenvatting

Diffusion Transformers hebben recentelijk opmerkelijke effectiviteit getoond bij het genereren van hoogwaardige 3D-puntenwolken. Het trainen van voxelgebaseerde diffusiemodellen voor hoogresolutie 3D-voxels blijft echter extreem kostbaar vanwege de kubische complexiteit van aandachtoperatoren, die voortkomt uit de extra dimensie van voxels. Gemotiveerd door de inherente redundantie van 3D in vergelijking met 2D, stellen we FastDiT-3D voor, een nieuwe gemaskerde diffusion transformer die is afgestemd op efficiënte 3D-puntenwolkengeneratie, waardoor de trainingskosten aanzienlijk worden verlaagd. Specifiek putten we inspiratie uit gemaskerde autoencoders om het denoisingsproces dynamisch uit te voeren op gemaskerde gevoxeliseerde puntenwolken. We introduceren ook een nieuwe voxelbewuste maskerstrategie om adaptief achtergrond-/voorgrondinformatie te aggregeren uit gevoxeliseerde puntenwolken. Onze methode behaalt state-of-the-art prestaties met een extreme maskeringsratio van bijna 99%. Bovendien introduceren we Mixture-of-Expert (MoE) in het 3D-diffusiemodel om de multi-categorie 3D-generatie te verbeteren. Elke categorie kan een apart diffusiepad leren met verschillende experts, waardoor gradientconflicten worden verlicht. Experimentele resultaten op de ShapeNet-dataset tonen aan dat onze methode state-of-the-art prestaties bereikt op het gebied van hoogwaardige en diverse 3D-puntenwolkengeneratie. Onze FastDiT-3D verbetert de 1-Nearest Neighbor Accuracy en Coverage-metrics bij het genereren van 128-resolutie voxel-puntenwolken, terwijl slechts 6,5% van de oorspronkelijke trainingskosten wordt gebruikt.

English

Diffusion Transformers have recently shown remarkable effectiveness in generating high-quality 3D point clouds. However, training voxel-based diffusion models for high-resolution 3D voxels remains prohibitively expensive due to the cubic complexity of attention operators, which arises from the additional dimension of voxels. Motivated by the inherent redundancy of 3D compared to 2D, we propose FastDiT-3D, a novel masked diffusion transformer tailored for efficient 3D point cloud generation, which greatly reduces training costs. Specifically, we draw inspiration from masked autoencoders to dynamically operate the denoising process on masked voxelized point clouds. We also propose a novel voxel-aware masking strategy to adaptively aggregate background/foreground information from voxelized point clouds. Our method achieves state-of-the-art performance with an extreme masking ratio of nearly 99%. Moreover, to improve multi-category 3D generation, we introduce Mixture-of-Expert (MoE) in 3D diffusion model. Each category can learn a distinct diffusion path with different experts, relieving gradient conflict. Experimental results on the ShapeNet dataset demonstrate that our method achieves state-of-the-art high-fidelity and diverse 3D point cloud generation performance. Our FastDiT-3D improves 1-Nearest Neighbor Accuracy and Coverage metrics when generating 128-resolution voxel point clouds, using only 6.5% of the original training cost.