Schnelles Training von Diffusionstransformatoren mit extremem Maskieren für die Generierung von 3D-Punktwolken

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

December 12, 2023

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

cs.AI

Zusammenfassung

Diffusion Transformers haben kürzlich bemerkenswerte Effektivität bei der Erzeugung hochwertiger 3D-Punktwolken gezeigt. Das Training von voxelbasierten Diffusionsmodellen für hochauflösende 3D-Voxel bleibt jedoch aufgrund der kubischen Komplexität der Aufmerksamkeitsoperatoren, die sich aus der zusätzlichen Dimension der Voxel ergibt, prohibitv teuer. Motiviert durch die inhärente Redundanz von 3D im Vergleich zu 2D, schlagen wir FastDiT-3D vor, einen neuartigen maskierten Diffusion Transformer, der für die effiziente Erzeugung von 3D-Punktwolken optimiert ist und die Trainingskosten erheblich reduziert. Insbesondere lassen wir uns von maskierten Autoencodern inspirieren, um den Denoising-Prozess dynamisch auf maskierten voxelisierten Punktwolken durchzuführen. Wir schlagen auch eine neuartige voxelbewusste Maskierungsstrategie vor, um Hintergrund-/Vordergrundinformationen aus voxelisierten Punktwolken adaptiv zu aggregieren. Unsere Methode erreicht state-of-the-art Leistung mit einem extremen Maskierungsverhältnis von fast 99%. Darüber hinaus führen wir Mixture-of-Expert (MoE) in 3D-Diffusionsmodellen ein, um die Erzeugung mehrerer Kategorien zu verbessern. Jede Kategorie kann einen eigenen Diffusionspfad mit verschiedenen Experten lernen, wodurch Gradientenkonflikte gelindert werden. Experimentelle Ergebnisse auf dem ShapeNet-Datensatz zeigen, dass unsere Methode state-of-the-art Leistung bei der Erzeugung hochwertiger und vielfältiger 3D-Punktwolken erreicht. Unser FastDiT-3D verbessert die 1-Nearest Neighbor Accuracy und Coverage Metriken bei der Erzeugung von 128-Auflösungs-Voxel-Punktwolken und verwendet dabei nur 6,5% der ursprünglichen Trainingskosten.

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.